From 7d91fa0890b4ce09dd5836fb48a403c563b8d5d1 Mon Sep 17 00:00:00 2001
From: erwei-xilinx <erwei.wang@amd.com>
Date: Mon, 24 Feb 2025 15:56:02 -0800
Subject: [PATCH] When lowering linalg to call, also link the parent air.herd
 to the func declaration

---
 mlir/lib/Conversion/AIRToAIEPass.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/lib/Conversion/AIRToAIEPass.cpp b/mlir/lib/Conversion/AIRToAIEPass.cpp
index dbcaeb641..a6e68fa2a 100644
--- a/mlir/lib/Conversion/AIRToAIEPass.cpp
+++ b/mlir/lib/Conversion/AIRToAIEPass.cpp
@@ -3679,6 +3679,12 @@ class AIRLinalgOpToLibraryCallRewrite
         op, fnNameAttr.getValue(), TypeRange(),
         ValueRange(ArrayRef<Value>(libFnOperands)));
 
+    if (auto herd = op->getParentOfType<air::HerdOp>())
+      rewriter.modifyOpInPlace(herd, [&]() {
+        herd->setAttr("link_with",
+                      StringAttr::get(rewriter.getContext(), linkWith));
+      });
+
     return success();
   }