Sketch out npu.configure op with ctrl pkt config

fixes

fixup tests

Move control packet data to end of runtime_sequence argument list

Author: fifield

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

@@ -907,6 +907,16 @@ def AIE_NpuControlPacketOp: AIEX_Op<"control_packet", []> {
   }];
 }
 
+def AIE_NpuConfigureOp: AIEX_Op<"npu.configure", []> {
+  let summary = "Configure NPU operation";
+  let arguments = (ins );
+  let results = (outs );
+  let assemblyFormat = [{ attr-dict }];
+  let description = [{
+    This operation is used to configure the NPU.
+  }];
+}
+
 // NPU Bd Write operation
 def AIE_NpuWriteBdOp: AIEX_Op<"npu.writebd", []> {
   let summary = "dma operator";

lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp

@@ -219,9 +219,7 @@ emitTransactionOps(OpBuilder &builder,
   auto loc = builder.getUnknownLoc();
 
   // create the txn ops
-  for (auto p : llvm::zip(operations, global_data)) {
-    auto op = std::get<0>(p);
-    memref::GlobalOp payload = std::get<1>(p);
+  for (auto [op, payload] : llvm::zip(operations, global_data)) {
 
     if (op.cmd.Opcode == XAie_TxnOpcode::XAIE_IO_WRITE) {
       builder.create<AIEX::NpuWrite32Op>(loc, op.cmd.RegOff, op.cmd.Value,
@@ -255,9 +253,7 @@ emitControlPacketOps(OpBuilder &builder,
   auto ctx = builder.getContext();
 
   // create the control packet ops
-  for (auto p : llvm::zip(operations, global_data)) {
-    auto op = std::get<0>(p);
-    memref::GlobalOp payload = std::get<1>(p);
+  for (auto [op, payload] : llvm::zip(operations, global_data)) {
 
     if (op.cmd.Opcode == XAie_TxnOpcode::XAIE_IO_WRITE) {
       builder.create<AIEX::NpuControlPacketOp>(
@@ -266,10 +262,10 @@ emitControlPacketOps(OpBuilder &builder,
           /*stream_id*/ builder.getI32IntegerAttr(0),
           DenseI32ArrayAttr::get(ctx, ArrayRef<int32_t>(op.cmd.Value)));
     } else if (op.cmd.Opcode == XAie_TxnOpcode::XAIE_IO_BLOCKWRITE) {
-      if (!std::get<1>(p).getInitialValue())
+      if (!payload.getInitialValue())
         continue;
       auto blockWriteData =
-          dyn_cast<DenseIntElementsAttr>(*std::get<1>(p).getInitialValue());
+          dyn_cast<DenseIntElementsAttr>(*payload.getInitialValue());
       if (!blockWriteData) {
         payload.emitError(
             "Global symbol initial value is not a dense int array");
@@ -383,30 +379,44 @@ static LogicalResult convertTransactionOpsToMLIR(
     global_data.push_back(global);
   }
 
-  // create aiex.runtime_sequence
-  int id = 0;
-  std::string seq_name = "configure";
-  while (device.lookupSymbol(seq_name))
-    seq_name = "configure" + std::to_string(id++);
-  StringAttr seq_sym_name = builder.getStringAttr(seq_name);
-  auto seq = builder.create<AIEX::RuntimeSequenceOp>(loc, seq_sym_name);
-  seq.getBody().push_back(new Block);
+  // search for npu.configure ops in runtime sequences by walking the device
+  // and collect them in a vector.
+  SmallVector<AIEX::NpuConfigureOp> configureOps;
+  device.walk([&](AIEX::NpuConfigureOp op) { configureOps.push_back(op); });
+
+  if (configureOps.empty()) {
+
+    // create aiex.runtime_sequence
+    int id = 0;
+    std::string seq_name = "configure";
+    while (device.lookupSymbol(seq_name))
+      seq_name = "configure" + std::to_string(id++);
+    StringAttr seq_sym_name = builder.getStringAttr(seq_name);
+    auto seq = builder.create<AIEX::RuntimeSequenceOp>(loc, seq_sym_name);
+    seq.getBody().push_back(new Block);
+
+    builder.setInsertionPointToStart(&seq.getBody().front());
+  } else {
+    builder.setInsertionPoint(configureOps.front());
+  }
 
   // create the txn ops
-  builder.setInsertionPointToStart(&seq.getBody().front());
   if (outputType == OutputType::Transaction) {
     if (failed(emitTransactionOps(builder, operations, global_data)))
       return failure();
   } else if (outputType == OutputType::ControlPacket) {
     if (failed(emitControlPacketOps(builder, operations, global_data)))
       return failure();
     // resolve mask writes; control packet doesn't natively support mask write.
-    if (failed(orConsecutiveWritesOnSameAddr(&seq.getBody().front())))
+    if (failed(orConsecutiveWritesOnSameAddr(builder.getBlock())))
       return failure();
   } else {
     llvm_unreachable("bad output type");
   }
 
+  if (!configureOps.empty())
+    configureOps.front().erase();
+
   return success();
 }
 

lib/Dialect/AIEX/Transforms/AIECtrlPacketToDma.cpp

@@ -14,6 +14,7 @@
 #include "aie/Dialect/AIEX/Transforms/AIEXPasses.h"
 
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/TypeSwitch.h"
@@ -72,21 +73,34 @@ struct AIECtrlPacketToDmaPass : AIECtrlPacketToDmaBase<AIECtrlPacketToDmaPass> {
 
       OpBuilder builder(f);
 
+      IRMapping mapping;
+
       auto newSeq =
           builder.create<AIEX::RuntimeSequenceOp>(loc, f.getSymNameAttr());
       newSeq.getBody().push_back(new Block);
 
+      // Copy the arguments from the old sequence to the new one.
+      for (auto arg : f.getBody().getArguments()) {
+        // Add the argument to the new sequence.
+        auto newArg = newSeq.getBody().addArgument(arg.getType(), arg.getLoc());
+        // Replace all uses of the old argument with the new one.
+        arg.replaceAllUsesWith(newArg);
+        // Add the mapping for the argument.
+        mapping.map(arg, newArg);
+      }
+
       // Using dynamic shape for ctrl pkt stream.
       auto ctrlPktMemrefType = MemRefType::get(
           ShapedType::kDynamic, IntegerType::get(ctx, 32), nullptr, 0);
       auto newBlockArg = newSeq.getBody().addArgument(ctrlPktMemrefType, loc);
+
       builder.setInsertionPointToStart(&newSeq.getBody().front());
 
       int64_t ddrOffset = 0;
       Block &entry = f.getBody().front();
       for (auto &o : entry) {
-        llvm::TypeSwitch<Operation *>(&o).Case<NpuControlPacketOp>(
-            [&](auto op) {
+        llvm::TypeSwitch<Operation *>(&o)
+            .Case<NpuControlPacketOp>([&](auto op) {
               // Destination tile info
               int col = op.getColumnFromAddr();
               int row = op.getRowFromAddr();
@@ -131,6 +145,10 @@ struct AIECtrlPacketToDmaPass : AIECtrlPacketToDmaBase<AIECtrlPacketToDmaPass> {
               auto row_num = builder.getI32IntegerAttr(1);
               builder.create<AIEX::NpuSyncOp>(loc, shimCol, shimRow, dir, chan,
                                               col_num, row_num);
+            })
+            .Default([&](Operation *op) {
+              // For all other operations, just clone them to the new sequence.
+              builder.clone(*op, mapping);
             });
       }
 

lib/Targets/AIETargetNPU.cpp

@@ -330,9 +330,11 @@ LogicalResult xilinx::AIE::AIETranslateControlPacketsToUI32Vec(
       int row = packetOp.getRowFromAddr();
       auto destTile = TileOp::getOrCreate(builder, deviceOp, col, row);
       auto info = destTile->getAttrOfType<AIE::PacketInfoAttr>("controller_id");
-      if (!info)
-        return destTile->emitError("Expected controller_id attribute");
-      uint32_t hdr = (info.getPktType() & 0x7) << 12 | (info.getPktId() & 0xff);
+      uint32_t hdr = 0;
+      if (info)
+        hdr = (info.getPktType() & 0x7) << 12 | (info.getPktId() & 0xff);
+      else
+        destTile->emitWarning("Expected controller_id attribute");
       words[0] = hdr | (0x1 & parity(hdr)) << 31;
 
       // control packet header

python/compiler/aiecc/main.py

@@ -31,6 +31,7 @@
 import aie.compiler.aiecc.cl_arguments
 import aie.compiler.aiecc.configure
 from aie.dialects import aie as aiedialect
+from aie.dialects import aiex as aiexdialect
 from aie.ir import Context, Location, Module
 from aie.passmanager import PassManager
 
@@ -614,7 +615,9 @@ async def process_txn(self, module_str):
                 print(f"copy {tmp} to {opts.txn_name}")
             shutil.copy(tmp, opts.txn_name)
 
-    async def aiebu_asm(self, input_file, output_file, ctrl_packet_file=None):
+    async def aiebu_asm(
+        self, input_file, output_file, ctrl_packet_file=None, ctrl_packet_idx=0
+    ):
 
         # find aiebu-asm binary
         asm_bin = "aiebu-asm"
@@ -645,7 +648,7 @@ async def aiebu_asm(self, input_file, output_file, ctrl_packet_file=None):
             exteral_buffers_json = {
                 "external_buffers": {
                     "buffer_ctrl": {
-                        "xrt_id": 0,
+                        "xrt_id": ctrl_packet_idx,
                         "logical_id": -1,
                         "size_in_bytes": ctrl_packet_size,
                         "ctrl_pkt_buffer": 1,
@@ -665,49 +668,53 @@ async def aiebu_asm(self, input_file, output_file, ctrl_packet_file=None):
         await self.do_call(None, args)
 
     async def process_ctrlpkt(self, module_str):
-        with Context(), Location.unknown():
-            run_passes(
-                "builtin.module(aie.device(convert-aie-to-control-packets{elf-dir="
-                + self.tmpdirname
-                + "}))",
-                module_str,
+        run_passes(
+            "builtin.module(aie.device(convert-aie-to-control-packets{elf-dir="
+            + self.tmpdirname
+            + "}))",
+            module_str,
+            self.prepend_tmp("ctrlpkt.mlir"),
+            self.opts.verbose,
+        )
+        await self.do_call(
+            None,
+            [
+                "aie-translate",
+                "-aie-ctrlpkt-to-bin",
                 self.prepend_tmp("ctrlpkt.mlir"),
-                self.opts.verbose,
-            )
-            await self.do_call(
-                None,
-                [
-                    "aie-translate",
-                    "-aie-ctrlpkt-to-bin",
-                    "-aie-sequence-name",
-                    "configure",
-                    self.prepend_tmp("ctrlpkt.mlir"),
-                    "-o",
-                    "ctrlpkt.bin",
-                ],
-            )
-            ctrlpkt_mlir_str = await read_file_async(self.prepend_tmp("ctrlpkt.mlir"))
-            run_passes(
-                "builtin.module(aie.device(aie-ctrl-packet-to-dma,aie-dma-to-npu))",
-                ctrlpkt_mlir_str,
+                "-o",
+                "ctrlpkt.bin",
+            ],
+        )
+        ctrlpkt_mlir_str = await read_file_async(self.prepend_tmp("ctrlpkt.mlir"))
+        run_passes(
+            "builtin.module(aie.device(aie-ctrl-packet-to-dma,aie-dma-to-npu))",
+            ctrlpkt_mlir_str,
+            self.prepend_tmp("ctrlpkt_dma_seq.mlir"),
+            self.opts.verbose,
+        )
+        await self.do_call(
+            None,
+            [
+                "aie-translate",
+                "-aie-npu-to-binary",
                 self.prepend_tmp("ctrlpkt_dma_seq.mlir"),
-                self.opts.verbose,
-            )
-            await self.do_call(
-                None,
-                [
-                    "aie-translate",
-                    "-aie-npu-to-binary",
-                    "-aie-sequence-name",
-                    "configure",
-                    self.prepend_tmp("ctrlpkt_dma_seq.mlir"),
-                    "-o",
-                    "ctrlpkt_dma_seq.bin",
-                ],
-            )
-            await self.aiebu_asm(
-                "ctrlpkt_dma_seq.bin", "ctrlpkt_dma_seq.elf", "ctrlpkt.bin"
+                "-o",
+                opts.insts_name,
+            ],
+        )
+        ctrl_idx = 0
+        ctrl_seq_str = await read_file_async(self.prepend_tmp("ctrlpkt_dma_seq.mlir"))
+        with Context(), Location.unknown():
+            dma_seq_module = Module.parse(ctrl_seq_str)
+            # walk through the dma sequence module to find runtime sequence
+            seqs = find_ops(
+                dma_seq_module.operation,
+                lambda o: isinstance(o.operation.opview, aiexdialect.RuntimeSequenceOp),
             )
+            if seqs:
+                ctrl_idx = len(seqs[0].regions[0].blocks[0].arguments.types) - 1
+        await self.aiebu_asm(opts.insts_name, opts.elf_name, "ctrlpkt.bin", ctrl_idx)
 
     async def process_elf(self, module_str):
         with Context(), Location.unknown():
@@ -1233,30 +1240,6 @@ async def run_flow(self):
                 exit(-3)
             aie_peano_target = aie_target.lower() + "-none-unknown-elf"
 
-            # Optionally generate insts.txt for NPU instruction stream
-            if opts.npu:
-                with Context(), Location.unknown():
-                    file_with_addresses_module = Module.parse(
-                        await read_file_async(file_with_addresses)
-                    )
-                    pass_pipeline = NPU_LOWERING_PIPELINE.materialize(module=True)
-                    npu_insts_file = (
-                        self.prepend_tmp("npu_insts.mlir")
-                        if self.opts.verbose
-                        else None
-                    )
-                    npu_insts_module = run_passes_module(
-                        pass_pipeline,
-                        file_with_addresses_module,
-                        npu_insts_file,
-                        self.opts.verbose,
-                    )
-                    npu_insts = aiedialect.translate_npu_to_binary(
-                        npu_insts_module.operation
-                    )
-                    with open(opts.insts_name, "wb") as f:
-                        f.write(struct.pack("I" * len(npu_insts), *npu_insts))
-
             # fmt: off
             if opts.unified:
                 file_opt_with_addresses = self.prepend_tmp("input_opt_with_addresses.mlir")
@@ -1349,6 +1332,30 @@ async def run_flow(self):
             if (opts.cdo or opts.xcl or opts.pdi) and opts.execute:
                 await self.process_cdo(input_physical_str)
 
+            # Optionally generate insts.txt for NPU instruction stream
+            if opts.npu and not opts.ctrlpkt:
+                with Context(), Location.unknown():
+                    file_with_addresses_module = Module.parse(
+                        await read_file_async(file_with_addresses)
+                    )
+                    pass_pipeline = NPU_LOWERING_PIPELINE.materialize(module=True)
+                    npu_insts_file = (
+                        self.prepend_tmp("npu_insts.mlir")
+                        if self.opts.verbose
+                        else None
+                    )
+                    npu_insts_module = run_passes_module(
+                        pass_pipeline,
+                        file_with_addresses_module,
+                        npu_insts_file,
+                        self.opts.verbose,
+                    )
+                    npu_insts = aiedialect.translate_npu_to_binary(
+                        npu_insts_module.operation
+                    )
+                    with open(opts.insts_name, "wb") as f:
+                        f.write(struct.pack("I" * len(npu_insts), *npu_insts))
+
             processes = []
             if opts.xcl:
                 processes.append(self.process_xclbin_gen())
@@ -1363,7 +1370,7 @@ async def run_flow(self):
             if opts.ctrlpkt and opts.execute:
                 processes.append(self.process_ctrlpkt(input_physical_str))
 
-            if opts.elf and opts.execute:
+            if opts.elf and not opts.ctrlpkt and opts.execute:
                 processes.append(self.process_elf(input_physical_str))
 
             await asyncio.gather(*processes)

test/npu-xrt/ctrl_packet_reconfig/aie2.mlir

@@ -67,12 +67,13 @@ module {
 
     aie.shim_dma_allocation @objFifo_in0(MM2S, 0, 0)
 
-    aiex.runtime_sequence @run(%arg0: memref<?xi8>, %arg1: memref<64x64xi8>, %arg2: memref<64x64xi8>) {
+    aiex.runtime_sequence @run(%arg1: memref<64x64xi8>, %arg2: memref<64x64xi8>) {
       %c0_i64 = arith.constant 0 : i64
       %c1_i64 = arith.constant 1 : i64
       %c56_i64 = arith.constant 56 : i64
       %c61_i64 = arith.constant 61 : i64
       %c64_i64 = arith.constant 64 : i64
+      aiex.npu.configure
       aiex.npu.dma_memcpy_nd (%arg1[%c0_i64, %c0_i64, %c0_i64, %c0_i64][%c1_i64, %c1_i64, %c64_i64, %c64_i64][%c0_i64, %c0_i64, %c64_i64, %c1_i64], packet = <pkt_id = 3, pkt_type = 0>) {id = 0 : i64, metadata = @objFifo_in0} : memref<64x64xi8>
       aiex.npu.dma_memcpy_nd (%arg2[%c0_i64, %c0_i64, %c0_i64, %c0_i64][%c1_i64, %c1_i64, %c64_i64, %c64_i64][%c0_i64, %c0_i64, %c64_i64, %c1_i64]) {id = 1 : i64, metadata = @objFifo_out0, issue_token = true} : memref<64x64xi8>
       aiex.npu.dma_wait { symbol = @objFifo_out0 }