diff --git a/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h b/mlir/include/air/Conversion/AIRToAIESchedulingUtils.h index d68f1a78f..96e45def5 100644 --- 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 index bf0a6899b..31639532e 100644 --- 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(SymbolTable::getSymbolAttrName()) + .getValue() + .str(); auto core = tile.getCoreOp(); if (!core) { core = builder.create(hloc, tile); tileToHerdMap[tile] = h; - auto herd_name = - aie_device - ->getAttrOfType(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(core_builder.getUnknownLoc(), + herd_lock, + AIE::LockAction::Acquire, 0); + } else if (aie_device.getTargetModel().getTargetArch() == + AIE::AIEArch::AIE2) { + core_builder.create( + 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(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(hloc, launch_bb); core_builder.setInsertionPoint(launch_bb->getTerminator()); - if (options.emit_herd_lock) - core_builder.create(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(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(entry_bb->getTerminator()); diff --git a/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp b/mlir/lib/Conversion/AIRToAIESchedulingUtils.cpp index d555dd9ff..1c9929c5f 100644 --- 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 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(t->getNextNode())) t = t->getNextNode(); b.setInsertionPointAfter(t); - return b.create(tile.getLoc(), tile, new_id, init); + auto lockOp = b.create(tile.getLoc(), tile, new_id, init); + if (name) + lockOp->setAttr(SymbolTable::getSymbolAttrName(), name); + return lockOp; } std::stringstream air::generateBufferNameInStringStream(StringRef prefix,