Enable support for drop-in replacement of load_pdi with control packets

Author: andrej

include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h

@@ -43,7 +43,8 @@ mlir::LogicalResult generateAndInsertConfigOps(
     llvm::StringRef clElfDir = "",
     AIEToConfigurationOutputType outputType =
         AIEToConfigurationOutputType::Transaction,
-    std::string blockwrite_prefix = "config_blockwrite_data_");
+    std::string blockwrite_prefix = "config_blockwrite_data_",
+    bool skipCtrlPktOverlay = false);
 
 // --------------------------------------------------------------------------
 // Device reset

include/aie/Dialect/AIE/Transforms/AIEPasses.h

@@ -59,9 +59,9 @@ std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIEAssignBufferDescriptorIDsPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIEObjectFifoLivenessPass();
-std::unique_ptr<mlir::OperationPass<DeviceOp>>
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 createAIEGenerateColumnControlOverlayPass();
-std::unique_ptr<mlir::OperationPass<DeviceOp>>
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 createAIEGenerateColumnControlOverlayPass(
     const AIEGenerateColumnControlOverlayOptions &options);
 std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIEAssignTileCtrlIDsPass();

include/aie/Dialect/AIE/Transforms/AIEPasses.td

@@ -408,7 +408,7 @@ def AIEAssignTileCtrlIDs : Pass<"aie-assign-tile-controller-ids", "DeviceOp"> {
   ];
 }
 
-def AIEGenerateColumnControlOverlay : Pass<"aie-generate-column-control-overlay", "DeviceOp"> {
+def AIEGenerateColumnControlOverlay : Pass<"aie-generate-column-control-overlay", "mlir::ModuleOp"> {
   let summary = "Spawns streaming interconnect network for column control";
   let description = [{
     For each column of AIE tiles being employed in the design, spawn a network of control packet
@@ -424,6 +424,10 @@ def AIEGenerateColumnControlOverlay : Pass<"aie-generate-column-control-overlay"
             "Flag to generate TCT routing between tile CTRL and shim SOUTH ports. Available options: ['shim-only', 'all-tiles', 'disable'].">,
     Option<"clRouteShimDmaToTileCTRL", "route-shim-to-tile-ctrl", "bool", /*default=*/"false",
             "Flag to generate routing between shim dma DMA and tile CTRL ports, for configuration.">,
+    Option<"clStandaloneDeviceOnly", "standalone-device-only", "bool", /*default=*/"false",
+            "Only emit the standalone @ctrl_pkt_overlay device; do not apply the overlay in place to source devices.">,
+    Option<"clEmitStandaloneOverlay", "emit-standalone-overlay", "bool", /*default=*/"false",
+            "Emit a standalone @ctrl_pkt_overlay device (and tag source devices) for the load-pdi-to-ctrl-pkt reconfigure flow. When false, the pass applies the overlay in place only (legacy behavior).">,
   ];
 }
 

include/aie/Dialect/AIEX/IR/AIEX.td

@@ -346,7 +346,8 @@ def AIE_ConfigureOp: AIEX_Op<"configure", [
  {
   let summary = "Set up a configuration (program memories, stream switches, etc.) on the NPU device.";
   let arguments = (
-    ins FlatSymbolRefAttr:$symbol
+    ins FlatSymbolRefAttr:$symbol,
+        OptionalAttr<AIEX_ExpandMode>:$expand_mode
   );
 
   let assemblyFormat = [{
@@ -977,7 +978,8 @@ def AIE_NpuLoadPdiOp: AIEX_Op<"npu.load_pdi", []> {
     ins OptionalAttr<FlatSymbolRefAttr>:$device_ref,
         DefaultValuedOptionalAttr<I32Attr, "0">:$id,
         DefaultValuedOptionalAttr<I32Attr, "0">:$size,
-        DefaultValuedOptionalAttr<UI64Attr, "0">:$address
+        DefaultValuedOptionalAttr<UI64Attr, "0">:$address,
+        OptionalAttr<AIEX_ExpandMode>:$expand_mode
   );
   let results = (outs );
   let assemblyFormat = [{ attr-dict }];

include/aie/Dialect/AIEX/IR/AIEXAttrs.td

@@ -33,4 +33,15 @@ def ScratchpadParameterKind : I32EnumAttr<"ScratchpadParameterKind",
   let cppNamespace = "::xilinx::AIEX";
 }
 
+// How a `load_pdi` op should be expanded by the `aie-expand-load-pdi` pass.
+def AIEX_ExpandNone    : I32EnumAttrCase<"none", 0>;
+def AIEX_ExpandWrite32 : I32EnumAttrCase<"write32", 1>;
+def AIEX_ExpandCtrlPkt : I32EnumAttrCase<"ctrlpkt", 2>;
+
+def AIEX_ExpandMode : I32EnumAttr<"ExpandMode",
+    "How to expand a load_pdi op",
+    [AIEX_ExpandNone, AIEX_ExpandWrite32, AIEX_ExpandCtrlPkt]> {
+  let cppNamespace = "::xilinx::AIEX";
+}
+
 #endif // AIEX_ATTRS

include/aie/Dialect/AIEX/Transforms/AIEXPasses.td

@@ -397,12 +397,27 @@ def AIELegalizeControlPacket : Pass<"aie-legalize-ctrl-packet", "AIE::DeviceOp">
 def AIEExpandLoadPdi : Pass<"aie-expand-load-pdi", "mlir::ModuleOp"> {
   let summary = "Expand load_pdi operations to explicit configuration sequences";
   let description = [{
-    This pass transforms `load_pdi { device_ref = @device }` operations as follows:
-    1. Create an empty device `@empty_N`.
-    2. Replace the load_pdi with `load_pdi { device_ref = @empty_N }` -- this causes a device reset.
-    3. Configure the device for the selected design by inserting write32s and blockwrites.
-
-    The configuration operations are generated by converting the referenced device's configuration to a transaction binary and then disassembling it into MLIR operations.
+    This pass transforms `load_pdi { device_ref = @device }` operations.
+
+    By default (`ctrl-pkt = false`), each op is replaced with:
+    1. A `load_pdi { device_ref = @empty_N }` op (loading an empty device,
+       which causes a device reset).
+    2. A sequence of `aiex.npu.write32`/`aiex.npu.blockwrite` ops that
+       configure the device.
+
+    With `ctrl-pkt = true`, each op is replaced with:
+    1. A `load_pdi { device_ref = @ctrl_pkt_overlay }` op, which loads the
+       column control packet overlay needed to stream further configuration
+       as control packets.
+    2. A sequence of `aiex.npu.control_packet` ops that configure the device.
+
+    The configuration operations are generated by converting the referenced
+    device's configuration to a transaction binary and then disassembling it
+    into MLIR operations.
+
+    This pass is intended to run after runtime sequence materialization
+    (i.e., after `aiex.configure` ops have been turned into
+    `aiex.npu.load_pdi`).
   }];
 
   let constructor = "xilinx::AIEX::createAIEExpandLoadPdiPass()";
@@ -411,6 +426,12 @@ def AIEExpandLoadPdi : Pass<"aie-expand-load-pdi", "mlir::ModuleOp"> {
         "xilinx::AIE::AIEDialect",
     "xilinx::AIEX::AIEXDialect",
   ];
+  let options = [
+    Option<"clCtrlPkt", "ctrl-pkt", "bool", /*default=*/"false",
+           "If true, lower each load_pdi to an @ctrl_pkt_overlay load_pdi "
+           "followed by aiex.npu.control_packet ops instead of write32/"
+           "blockwrite ops with an empty-device reset.">,
+  ];
 }
 
 def AIEXInlineTraceConfig : Pass<"aie-inline-trace-config", "AIE::DeviceOp"> {

include/aie/Targets/AIERT.h

@@ -55,8 +55,10 @@ struct AIERTControl {
   mlir::LogicalResult configureLocksAndBd(mlir::Block &block, int col, int row);
   mlir::LogicalResult initLocks(DeviceOp &targetOp);
   mlir::LogicalResult initBuffers(DeviceOp &targetOp);
-  mlir::LogicalResult configureSwitches(DeviceOp &targetOp);
-  mlir::LogicalResult addInitConfig(DeviceOp &targetOp);
+  mlir::LogicalResult configureSwitches(DeviceOp &targetOp,
+                                        bool skipCtrlPktOverlay = false);
+  mlir::LogicalResult addInitConfig(DeviceOp &targetOp,
+                                    bool skipCtrlPktOverlay = false);
   mlir::LogicalResult addCoreEnable(DeviceOp &targetOp);
   mlir::LogicalResult addAieElf(uint8_t col, uint8_t row,
                                 const mlir::StringRef elfPath, bool aieSim);

lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp

@@ -436,11 +436,12 @@ static LogicalResult generateTransactions(AIERTControl &ctl,
                                           const StringRef workDirPath,
                                           DeviceOp &targetOp, bool aieSim,
                                           bool enableElfs, bool enableInit,
-                                          bool enableCores) {
+                                          bool enableCores,
+                                          bool skipCtrlPktOverlay = false) {
   if (enableElfs && !targetOp.getOps<CoreOp>().empty() &&
       failed(ctl.addAieElfs(targetOp, workDirPath, aieSim)))
     return failure();
-  if (enableInit && failed(ctl.addInitConfig(targetOp)))
+  if (enableInit && failed(ctl.addInitConfig(targetOp, skipCtrlPktOverlay)))
     return failure();
   if (enableCores && !targetOp.getOps<CoreOp>().empty() &&
       failed(ctl.addCoreEnable(targetOp)))
@@ -508,8 +509,9 @@ emitTransactionOps(OpBuilder &builder, Location fallbackLoc,
       IntegerAttr addressAttr =
           IntegerAttr::get(ui64Ty, llvm::APInt(64, payloadInfo.address));
 
-      AIEX::NpuLoadPdiOp::create(builder, loc, nullptr, idAttr, sizeAttr,
-                                 addressAttr);
+      AIEX::NpuLoadPdiOp::create(builder, loc, /*device_ref=*/nullptr, idAttr,
+                                 sizeAttr, addressAttr,
+                                 /*expand_mode=*/nullptr);
     } else if (op.cmd.Opcode == XAie_TxnOpcode::XAIE_IO_CUSTOM_OP_DDR_PATCH) {
       if (!op.addressPatch) {
         llvm::errs()
@@ -770,7 +772,7 @@ xilinx::AIE::convertTransactionBinaryToMLIR(mlir::MLIRContext *ctx,
 LogicalResult xilinx::AIE::generateAndInsertConfigOps(
     OpBuilder &builder, xilinx::AIE::DeviceOp device, llvm::StringRef clElfDir,
     AIE::AIEToConfigurationOutputType outputType,
-    std::string blockwrite_prefix) {
+    std::string blockwrite_prefix, bool skipCtrlPktOverlay) {
   const AIETargetModel &targetModel =
       (const AIETargetModel &)device.getTargetModel();
 
@@ -789,7 +791,7 @@ LogicalResult xilinx::AIE::generateAndInsertConfigOps(
 
   bool generateElfs = true;
   if (failed(generateTransactions(ctl, clElfDir, device, aieSim, generateElfs,
-                                  true, true)))
+                                  true, true, skipCtrlPktOverlay)))
     return failure();
 
   // Export the transactions to a binary buffer

lib/Dialect/AIE/Transforms/AIECreatePathFindFlows.cpp

@@ -282,6 +282,8 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
   DenseMap<PhysPort, BoolAttr> keepPktHeaderAttr;
   // Map from tileID and master ports to flags labelling control packet flows
   DenseMap<std::pair<PhysPort, int>, bool> ctrlPktFlows;
+  // Set of master ports that belong to control packet overlay flows
+  DenseSet<PhysPort> ctrlPktOverlayMasterPorts;
 
   for (auto tileOp : device.getOps<TileOp>()) {
     int col = tileOp.colIndex();
@@ -365,9 +367,10 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
       Port destPort = conn.dst;
       auto sourceFlow =
           std::make_pair(std::make_pair(tileId, sourcePort), flowID);
-      if (ctrlPktFlows[{{tileId, destPort}, flowID}])
+      if (ctrlPktFlows[{{tileId, destPort}, flowID}]) {
         ctrlPacketFlows[sourceFlow].push_back({tileId, destPort});
-      else
+        ctrlPktOverlayMasterPorts.insert({tileId, destPort});
+      } else
         packetFlows[sourceFlow].push_back({tileId, destPort});
       slavePorts.push_back(sourceFlow);
       LLVM_DEBUG(llvm::dbgs() << "flowID " << flowID << ':'
@@ -882,9 +885,11 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
         amsels.push_back(amselOps[msel]);
       }
 
-      MasterSetOp::create(builder, tileLoc, builder.getIndexType(), bundle,
-                          channel, amsels,
-                          keepPktHeaderAttr[{tileId, tileMaster}]);
+      auto msOp = MasterSetOp::create(builder, tileLoc, builder.getIndexType(),
+                                      bundle, channel, amsels,
+                                      keepPktHeaderAttr[{tileId, tileMaster}]);
+      if (ctrlPktOverlayMasterPorts.contains({tileId, tileMaster}))
+        msOp->setAttr("is_ctrl_pkt_overlay", builder.getUnitAttr());
     }
 
     // Generate the packet rules
@@ -910,11 +915,16 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
 #endif
       Value amsel = amselOps[slaveAMSels[group.front()]];
 
+      // Check if this group is a ctrl-pkt overlay flow
+      bool isCtrlPktGroup = ctrlPacketFlows.count(group.front()) > 0;
+
       PacketRulesOp packetrules;
       if (slaveRules.count(slave) == 0) {
         packetrules = PacketRulesOp::create(builder, tileLoc, bundle, channel);
         PacketRulesOp::ensureTerminator(packetrules.getRules(), builder,
                                         tileLoc);
+        if (isCtrlPktGroup)
+          packetrules->setAttr("is_ctrl_pkt_overlay", builder.getUnitAttr());
         slaveRules[slave] = packetrules;
       } else
         packetrules = slaveRules[slave];