Configure AIE using control packets

include/aie-c/Translation.h

@@ -19,6 +19,7 @@ MLIR_CAPI_EXPORTED MlirStringRef aieTranslateAIEVecToCpp(MlirOperation op,
                                                          bool aie2);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateModuleToLLVMIR(MlirOperation op);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToNPU(MlirOperation op);
+MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToControlPackets(MlirOperation op);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToXAIEV2(MlirOperation op);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToHSA(MlirOperation op);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToBCF(MlirOperation op, int col,

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

@@ -791,6 +791,27 @@ def AIE_NpuAddressPatchOp: AIEX_Op<"npu.address_patch", []> {
   }];
 }
 
+def AIE_NpuControlPacketOp: AIEX_Op<"npu.control_packet", []> {
+  let summary = "AIE control packet";
+  let arguments = (
+    ins UI32Attr:$address,
+        OptionalAttr<I32Attr>:$length,
+        I32Attr:$opcode,
+        I32Attr:$stream_id,
+        OptionalAttr<DenseI32ArrayAttr>:$data
+  );
+  let results = (outs );
+  let assemblyFormat = [{ attr-dict }];
+  let description = [{
+    The control_packet operation represents a low-level AIE control packet header
+    and payload.
+  }];
+  let extraClassDeclaration = [{
+    uint32_t getRow();
+    uint32_t getColumn();
+  }];
+}
+
 // NPU Bd Write operation
 def AIE_NpuWriteBdOp: AIEX_Op<"npu.writebd", []> {
   let summary = "dma operator";

include/aie/Targets/AIETargets.h

@@ -33,7 +33,11 @@ mlir::LogicalResult AIETranslateGraphXPE(mlir::ModuleOp module,
                                          llvm::raw_ostream &);
 mlir::LogicalResult AIETranslateToNPU(mlir::ModuleOp module,
                                       llvm::raw_ostream &output);
-std::vector<uint32_t> AIETranslateToNPU(mlir::ModuleOp);
+mlir::LogicalResult AIETranslateToNPU(mlir::ModuleOp, std::vector<uint32_t> &);
+mlir::LogicalResult AIETranslateToControlPackets(mlir::ModuleOp module,
+                                                 llvm::raw_ostream &output);
+mlir::LogicalResult AIETranslateToControlPackets(mlir::ModuleOp,
+                                                 std::vector<uint32_t> &);
 mlir::LogicalResult AIETranslateToLdScript(mlir::ModuleOp module,
                                            llvm::raw_ostream &output,
                                            int tileCol, int tileRow);

lib/CAPI/Translation.cpp

@@ -109,6 +109,17 @@ MlirStringRef aieTranslateToNPU(MlirOperation moduleOp) {
   return mlirStringRefCreate(cStr, npu.size());
 }
 
+MlirStringRef aieTranslateToControlPackets(MlirOperation moduleOp) {
+  std::string npu;
+  llvm::raw_string_ostream os(npu);
+  ModuleOp mod = llvm::cast<ModuleOp>(unwrap(moduleOp));
+  if (failed(AIETranslateToControlPackets(mod, os)))
+    return mlirStringRefCreate(nullptr, 0);
+  char *cStr = static_cast<char *>(malloc(npu.size()));
+  npu.copy(cStr, npu.size());
+  return mlirStringRefCreate(cStr, npu.size());
+}
+
 MlirStringRef aieTranslateToXAIEV2(MlirOperation moduleOp) {
   std::string xaie;
   llvm::raw_string_ostream os(xaie);

lib/Dialect/AIEX/IR/AIEXDialect.cpp

@@ -227,13 +227,13 @@
     return forOp->emitOpError(
         "Size 1 exceeds the [0:" + std::to_string((1 << wrap_bits) - 1) +
         "] range.");
   if (hardwareSizes[3] > (1 << iter_bits))
     return forOp->emitOpError(
         "Size 3 exceeds the [1:" + std::to_string(1 << iter_bits) + "] range.");
   if (hardwareStrides[0] > (1 << step_bits))
     return forOp->emitOpError("Stride 0 exceeds the [1:" +
                               std::to_string(1 << step_bits) + "] range.");
   if (hardwareStrides[1] > (1 << step_bits))
     return forOp->emitOpError("Stride 1 exceeds the [1:" +
                               std::to_string(1 << step_bits) + "] range.");
   if (hardwareStrides[2] > (1 << step_bits))
@@ -641,4 +641,24 @@
     }
   }
   return success();
-}
+}
+
+//===----------------------------------------------------------------------===//
+// NpuControlPacketOp
+//===----------------------------------------------------------------------===//
+
+uint32_t AIEX::NpuControlPacketOp::getRow() {
+  const auto &targetModel = AIE::getTargetModel(*this);
+  uint32_t addr = getAddress();
+  uint32_t rowInt =
+      (addr & (0xff << targetModel.getRowShift())) >> targetModel.getRowShift();
+  return rowInt;
+}
+
+uint32_t AIEX::NpuControlPacketOp::getColumn() {
+  const auto &targetModel = AIE::getTargetModel(*this);
+  uint32_t addr = getAddress();
+  uint32_t colInt = (addr & (0xff << targetModel.getColumnShift())) >>
+                    targetModel.getColumnShift();
+  return colInt;
+}

lib/Targets/AIETargetNPU.cpp

@@ -185,9 +185,9 @@ void appendBlockWrite(std::vector<uint32_t> &instructions, NpuBlockWriteOp op) {
 
 } // namespace
 
-std::vector<uint32_t> xilinx::AIE::AIETranslateToNPU(ModuleOp module) {
-
-  std::vector<uint32_t> instructions;
+LogicalResult
+xilinx::AIE::AIETranslateToNPU(ModuleOp module,
+                               std::vector<uint32_t> &instructions) {
 
   auto words = reserveAndGetTail(instructions, 4);
 
@@ -236,12 +236,67 @@ std::vector<uint32_t> xilinx::AIE::AIETranslateToNPU(ModuleOp module) {
   // write size fields of the txn header
   instructions[2] = count;
   instructions[3] = instructions.size() * sizeof(uint32_t); // size of the txn
-  return instructions;
+  return success();
 }
 
 LogicalResult xilinx::AIE::AIETranslateToNPU(ModuleOp module,
                                              raw_ostream &output) {
-  auto instructions = AIETranslateToNPU(module);
+  std::vector<uint32_t> instructions;
+  auto r = AIETranslateToNPU(module, instructions);
+  if (failed(r))
+    return r;
+  for (auto w : instructions)
+    output << llvm::format("%08X\n", w);
+  return success();
+}
+
+LogicalResult
+xilinx::AIE::AIETranslateToControlPackets(ModuleOp module,
+                                          std::vector<uint32_t> &instructions) {
+
+  DeviceOp deviceOp = *module.getOps<DeviceOp>().begin();
+  auto sequenceOps = deviceOp.getOps<AIEX::RuntimeSequenceOp>();
+  for (auto f : sequenceOps) {
+    Block &entry = f.getBody().front();
+    for (auto &o : entry) {
+      llvm::TypeSwitch<Operation *>(&o).Case<NpuControlPacketOp>([&](auto op) {
+        uint32_t size = 0;
+        auto data = op.getData();
+        auto length = op.getLength();
+        if (data)
+          size = data->size();
+        auto words = reserveAndGetTail(instructions, 1 + size);
+        if (!data && length)
+          size = *length;
+        auto parity = [](uint32_t n) {
+          uint32_t p = 0;
+          while (n) {
+            p += n & 1;
+            n >>= 1;
+          }
+          return (p % 2) == 0;
+        };
+        uint32_t addr = op.getAddress() & 0xFFFFF;
+        uint32_t beats = size - 1;
+        uint32_t opc = op.getOpcode();
+        uint32_t id = op.getStreamId();
+        uint32_t hdr = id << 24 | opc << 22 | beats << 20 | addr;
+        words[0] = hdr | (0x1 & parity(hdr)) << 31;
+        if (opc == 0x0 || opc == 0x2)
+          for (unsigned i = 0; i < size; i++)
+            words[i + 1] = data.value()[i];
+      });
+    }
+  }
+  return success();
+}
+
+LogicalResult xilinx::AIE::AIETranslateToControlPackets(ModuleOp module,
+                                                        raw_ostream &output) {
+  std::vector<uint32_t> instructions;
+  auto r = AIETranslateToControlPackets(module, instructions);
+  if (failed(r))
+    return r;
   for (auto w : instructions)
     output << llvm::format("%08X\n", w);
   return success();

lib/Targets/AIETargets.cpp

@@ -153,9 +153,11 @@ void registerAIETranslations() {
       "cdo-enable-cores", llvm::cl::init(true),
       llvm::cl::desc("Enable cores in CDO"));
 
-  static llvm::cl::opt<bool> npuInstGenBinary(
-      "aie-npu-instgen-binary", llvm::cl::init(false),
-      llvm::cl::desc("Emit binary (true) or text (false) NPU instructions"));
+  static llvm::cl::opt<bool> outputBinary(
+      "aie-output-binary", llvm::cl::init(false),
+      llvm::cl::desc(
+          "Select binary (true) or text (false) output for supported "
+          "translations. e.g. aie-npu-instgen, aie-ctrlpkt-to-bin"));
 
   TranslateFromMLIRRegistration registrationMMap(
       "aie-generate-mmap", "Generate AIE memory map",
@@ -352,16 +354,34 @@ void registerAIETranslations() {
       },
       registerDialects);
   TranslateFromMLIRRegistration registrationNPU(
-      "aie-npu-instgen", "Generate instructions for NPU",
+      "aie-npu-instgen", "Translate npu instructions to binary",
       [](ModuleOp module, raw_ostream &output) {
-        if (npuInstGenBinary == true) {
-          auto instructions = AIETranslateToNPU(module);
+        if (outputBinary == true) {
+          std::vector<uint32_t> instructions;
+          auto r = AIETranslateToNPU(module, instructions);
+          if (failed(r))
+            return r;
           output.write(reinterpret_cast<const char *>(instructions.data()),
                        instructions.size() * sizeof(uint32_t));
           return success();
         }
         return AIETranslateToNPU(module, output);
       },
       registerDialects);
+  TranslateFromMLIRRegistration registrationCtrlPkt(
+      "aie-ctrlpkt-to-bin", "Translate aiex.control_packet ops to binary",
+      [](ModuleOp module, raw_ostream &output) {
+        if (outputBinary == true) {
+          std::vector<uint32_t> instructions;
+          auto r = AIETranslateToControlPackets(module, instructions);
+          if (failed(r))
+            return r;
+          output.write(reinterpret_cast<const char *>(instructions.data()),
+                       instructions.size() * sizeof(uint32_t));
+          return success();
+        }
+        return AIETranslateToControlPackets(module, output);
+      },
+      registerDialects);
 }
 } // namespace xilinx::AIE

python/AIEMLIRModule.cpp

@@ -142,6 +142,18 @@ PYBIND11_MODULE(_aie, m) {
       },
       "module"_a);
 
+  m.def(
+      "generate_control_packets",
+      [&stealCStr](MlirOperation op) {
+        py::str ctrlPackets = stealCStr(aieTranslateToControlPackets(op));
+        auto individualInstructions =
+            ctrlPackets.attr("split")().cast<py::list>();
+        for (size_t i = 0; i < individualInstructions.size(); ++i)
+          individualInstructions[i] = individualInstructions[i].attr("strip")();
+        return individualInstructions;
+      },
+      "module"_a);
+
   m.def(
       "generate_xaie",
       [&stealCStr](MlirOperation op) {