Add npu.load_pdi to aiex (#2147)

Author: fifield

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

@@ -953,6 +953,23 @@ def AIE_NpuWriteBdOp: AIEX_Op<"npu.writebd", []> {
   }];
 }
 
+// NPU Load PDI operation
+def AIE_NpuLoadPdiOp: AIEX_Op<"npu.load_pdi", []> {
+  let summary = "load pdi operator";
+  let arguments = (
+    ins I32Attr:$id,
+        OptionalAttr<I32Attr>:$size,
+        OptionalAttr<UI64Attr>:$address
+  );
+  let results = (outs );
+  let assemblyFormat = [{ attr-dict }];
+  let description = [{
+    Load a PDI (Programmable Device Image) to configure the NPU.
+    The PDI is identified by `id`. `address` and `size` are typically written at
+    runtime by the driver or host program.
+  }];
+}
+
 def AIE_DMAConfigureTaskOp : AIEX_Op<"dma_configure_task", [HasParent<"RuntimeSequenceOp">, TileElement]>, Results<(outs Index:$result)> {
   let summary = "Concrete Instantiation of a Buffer Descriptor Chain as a Task on a Channel and Direction on a Tile";
   let description = [{

lib/Targets/AIETargetNPU.cpp

@@ -148,6 +148,23 @@ void appendMaskWrite32(std::vector<uint32_t> &instructions,
   words[6] = words.size() * sizeof(uint32_t); // Operation Size
 }
 
+void appendLoadPdi(std::vector<uint32_t> &instructions, NpuLoadPdiOp op) {
+
+  auto words = reserveAndGetTail(instructions, 4);
+
+  // XAIE_IO_LOADPDI
+  words[0] = XAIE_IO_LOADPDI;
+  words[0] |= op.getId() << 16;
+  std::optional<uint32_t> size = op.getSize();
+  if (size)
+    words[1] = *size;
+  std::optional<uint64_t> address = op.getAddress();
+  if (address) {
+    words[2] = *address;
+    words[3] = *address >> 32;
+  }
+}
+
 void appendAddressPatch(std::vector<uint32_t> &instructions,
                         NpuAddressPatchOp op) {
 
@@ -277,6 +294,10 @@ xilinx::AIE::AIETranslateNpuToBinary(ModuleOp module,
             count++;
             appendMaskWrite32(instructions, op);
           })
+          .Case<NpuLoadPdiOp>([&](auto op) {
+            count++;
+            appendLoadPdi(instructions, op);
+          })
           .Case<NpuAddressPatchOp>([&](auto op) {
             count++;
             appendAddressPatch(instructions, op);

test/Targets/NPU/npu_instgen.mlir

@@ -17,8 +17,8 @@ module {
       // TXN header 0.1
       // CHECK: 06030100
       // CHECK: 00000104
-      // CHECK: 00000006
-      // CHECK: 000000CC
+      // CHECK: 00000007
+      // CHECK: 000000DC
 
       // CHECK: 00000000
       // CHECK: 00000000
@@ -79,6 +79,12 @@ module {
       // CHECK: 00030401
       // CHECK: 05010200
       aiex.npu.sync { column = 3 : i32, row = 4 : i32, direction = 1 : i32, channel = 5 : i32, column_num = 1 : i32, row_num = 2 : i32 }
+
+      // CHECK: 00020008
+      // CHECK: 00000400
+      // CHECK: 12345678
+      // CHECK: 00000ABC
+      aiex.npu.load_pdi { address = 0xabc12345678 : ui64, id = 2 : i32, size = 1024 : i32 }
     }
   }
 }

test/dialect/AIEX/roundtrip.mlir

@@ -64,3 +64,19 @@ aie.device(npu1) {
     aiex.control_packet {address = 0xABCD : ui32, length = 3 : i32, opcode = 1 : i32, stream_id = 4 : i32}
   }
 }
+
+// -----
+
+// CHECK: aie.device
+// CHECK: aiex.npu.load_pdi {id = 4 : i32}
+// CHECK: aiex.npu.load_pdi {id = 7 : i32}
+// CHECK: aiex.npu.load_pdi {address = 2 : ui64, id = 1 : i32, size = 3 : i32}
+aie.device(npu1) {
+  memref.global "private" constant @pdi_data : memref<8xi32> = dense<[-1, 1, 256, 1024, 41, 42, 43, 44]>
+  aiex.runtime_sequence @pdi_loader(%arg0 : memref<?xi32>) {
+    %0 = memref.get_global @pdi_data : memref<8xi32>
+    aiex.npu.load_pdi {id = 4 : i32}
+    aiex.npu.load_pdi {id = 7 : i32}
+    aiex.npu.load_pdi {id = 1 : i32, address = 2 : ui64, size = 3 : i32}
+  }
+}