update herd lock for aie2

Xilinx · Jun 21, 2024 · 318d6ed · 318d6ed
1 parent 568d409
commit 318d6ed
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 19 deletions.
diff --git a/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h b/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h
@@ -27,7 +27,7 @@ AIE::TileOp getPhysTileOpOrNull(AIE::DeviceOp aie_device, int col, int row);
 AIE::TileOp getPhysTileOp(AIE::DeviceOp aie_device, int col, int row);
 
 AIE::LockOp allocateLockOp(AIE::DeviceOp aie_device, AIE::TileOp tile,
-                           int init = 0, int id = -1);
+                           int init = 0, int id = -1, StringAttr name = nullptr);
 
 std::stringstream
 generateBufferNameInStringStream(StringRef prefix, uint64_t &BufferId,

diff --git a/mlir/lib/Conversion/AIRToAIEPass.cpp b/mlir/lib/Conversion/AIRToAIEPass.cpp
@@ -244,15 +244,15 @@ void outlineAIECores(OpBuilder &builder, AIE::DeviceOp aie_device,
       builder.setInsertionPointAfter(t);
 
       // make the aie.core for the tile core
+      auto herd_name =
+          aie_device
+              ->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
+              .getValue()
+              .str();
       auto core = tile.getCoreOp();
       if (!core) {
         core = builder.create<AIE::CoreOp>(hloc, tile);
         tileToHerdMap[tile] = h;
-        auto herd_name =
-            aie_device
-                ->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
-                .getValue()
-                .str();
         core->setAttr(
             "elf_file",
             StringAttr::get(ctx, herd_name + "_core_" + std::to_string(phys_x) +
@@ -283,8 +283,14 @@ void outlineAIECores(OpBuilder &builder, AIE::DeviceOp aie_device,
       }
 
       Value herd_lock = nullptr;
-      if (options.emit_herd_lock)
-        herd_lock = allocateLockOp(aie_device, tile, /*init=*/0, /*id=*/0);
+      if (options.emit_herd_lock) {
+        StringAttr name =
+            builder.getStringAttr("__air_herd_lock_" + std::to_string(phys_x) +
+                                  "_" + std::to_string(phys_y));
+        // herd lock is always lock zero
+        herd_lock =
+            allocateLockOp(aie_device, tile, /*init=*/0, /*id=*/0, name);
+      }
 
       assert((h.getBody().getBlocks().size() == 1) &&
              "Launch body can only contain one Block");
@@ -346,6 +352,18 @@ void outlineAIECores(OpBuilder &builder, AIE::DeviceOp aie_device,
               AIE::LockAction::AcquireGreaterEqual, 1);
         }
       }
+      if (options.emit_herd_lock) {
+        if (aie_device.getTargetModel().getTargetArch() == AIE::AIEArch::AIE1) {
+          core_builder.create<AIE::UseLockOp>(core_builder.getUnknownLoc(),
+                                              herd_lock,
+                                              AIE::LockAction::Acquire, 0);
+        } else if (aie_device.getTargetModel().getTargetArch() ==
+                   AIE::AIEArch::AIE2) {
+          core_builder.create<AIE::UseLockOp>(
+              core_builder.getUnknownLoc(), herd_lock,
+              AIE::LockAction::AcquireGreaterEqual, 1);
+        }
+      }
 
       for (unsigned ki = 0, ke = h.getNumKernelOperands(); ki < ke; ki++) {
         BlockArgument karg = h.getKernelArgument(ki);
@@ -418,21 +436,23 @@ void outlineAIECores(OpBuilder &builder, AIE::DeviceOp aie_device,
         remap.map(karg, m);
       }
 
-      if (options.emit_herd_lock)
-        core_builder.create<AIE::UseLockOp>(core_builder.getUnknownLoc(),
-                                            herd_lock, AIE::LockAction::Acquire,
-                                            0);
-
       Region &r = h.getRegion();
       r.cloneInto(&core.getBody(), remap);
 
       Block *launch_bb = remap.lookup(&r.front());
       core_builder.create<cf::BranchOp>(hloc, launch_bb);
       core_builder.setInsertionPoint(launch_bb->getTerminator());
-      if (options.emit_herd_lock)
-        core_builder.create<AIE::UseLockOp>(core_builder.getUnknownLoc(),
-                                            herd_lock, AIE::LockAction::Release,
-                                            0);
+      if (options.emit_herd_lock) {
+        if (aie_device.getTargetModel().getTargetArch() == AIE::AIEArch::AIE1) {
+          core_builder.create<AIE::UseLockOp>(core_builder.getUnknownLoc(),
+                                              herd_lock,
+                                              AIE::LockAction::Release, 0);
+        } else if (aie_device.getTargetModel().getTargetArch() ==
+                   AIE::AIEArch::AIE2) {
+          // we could release something, but we don't have to a way to observe
+          // it yet in NPU
+        }
+      }
 
       if (options.emit_while) {
         auto entry_bb_br = dyn_cast<cf::BranchOp>(entry_bb->getTerminator());

diff --git a/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp b/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp
@@ -77,7 +77,7 @@ AIE::TileOp air::getPhysTileOp(AIE::DeviceOp aie_device, int col, int row) {
 }
 
 AIE::LockOp air::allocateLockOp(AIE::DeviceOp aie_device, AIE::TileOp tile,
-                                int init, int id) {
+                                int init, int id, StringAttr name) {
   AIE::LockOp lock = nullptr;
   std::set<int> ids;
   aie_device.walk([&](AIE::LockOp l) {
@@ -105,7 +105,10 @@ AIE::LockOp air::allocateLockOp(AIE::DeviceOp aie_device, AIE::TileOp tile,
   while (dyn_cast_or_null<AIE::TileOp>(t->getNextNode()))
     t = t->getNextNode();
   b.setInsertionPointAfter(t);
-  return b.create<AIE::LockOp>(tile.getLoc(), tile, new_id, init);
+  auto lockOp = b.create<AIE::LockOp>(tile.getLoc(), tile, new_id, init);
+  if (name)
+    lockOp->setAttr(SymbolTable::getSymbolAttrName(), name);
+  return lockOp;
 }
 
 std::stringstream air::generateBufferNameInStringStream(StringRef prefix,