Add command line device selection

Author: fifield

include/aie-c/Translation.h

@@ -40,8 +40,8 @@ MLIR_CAPI_EXPORTED MlirLogicalResult aieTranslateToCDODirect(
     bool xaieDebug, bool enableCores);
 MLIR_CAPI_EXPORTED MlirOperation aieTranslateBinaryToTxn(MlirContext ctx,
                                                          MlirStringRef binary);
-MLIR_CAPI_EXPORTED MlirOperation
-aieTranslateBinaryToControlPackets(MlirContext ctx, MlirStringRef binary);
+MLIR_CAPI_EXPORTED MlirOperation aieTranslateBinaryToControlPackets(
+    MlirContext ctx, MlirStringRef binary, int device);
 
 struct AieRtControl {
   void *ptr;

include/aie/Conversion/AIEToConfiguration/AIEToConfiguration.h

@@ -17,6 +17,7 @@
 
 namespace xilinx::AIE {
 
+enum class AIEDevice : uint32_t;
 class DeviceOp;
 
 std::unique_ptr<mlir::OperationPass<xilinx::AIE::DeviceOp>>
@@ -29,9 +30,9 @@ std::optional<mlir::ModuleOp>
 convertTransactionBinaryToMLIR(mlir::MLIRContext *ctx,
                                std::vector<uint8_t> &binary);
 
-std::optional<mlir::ModuleOp>
-convertControlPacketBinaryToMLIR(mlir::MLIRContext *ctx,
-                                 std::vector<uint8_t> &binary);
+std::optional<mlir::ModuleOp> convertControlPacketBinaryToMLIR(
+    mlir::MLIRContext *ctx, std::vector<uint8_t> &binary,
+    AIEDevice device = static_cast<AIEDevice>(4)); /* npu1 */
 
 } // namespace xilinx::AIE
 

lib/CAPI/Translation.cpp

@@ -95,9 +95,11 @@ MlirOperation aieTranslateBinaryToTxn(MlirContext ctx, MlirStringRef binary) {
 }
 
 MlirOperation aieTranslateBinaryToControlPackets(MlirContext ctx,
-                                                 MlirStringRef binary) {
+                                                 MlirStringRef binary,
+                                                 int device) {
   std::vector<uint8_t> binaryData(binary.data, binary.data + binary.length);
-  auto mod = convertControlPacketBinaryToMLIR(unwrap(ctx), binaryData);
+  auto mod = convertControlPacketBinaryToMLIR(unwrap(ctx), binaryData,
+                                              static_cast<AIEDevice>(device));
   if (!mod)
     return wrap(ModuleOp().getOperation());
   return wrap(mod->getOperation());

lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp

@@ -863,9 +863,8 @@ xilinx::AIE::convertTransactionBinaryToMLIR(mlir::MLIRContext *ctx,
 // new ModuleOp containing a DeviceOp containing a runtime sequence with the
 // control packet binary encoded as a sequence of aiex.control_packet
 // operations. On failure return std::nullopt.
-std::optional<mlir::ModuleOp>
-xilinx::AIE::convertControlPacketBinaryToMLIR(mlir::MLIRContext *ctx,
-                                              std::vector<uint8_t> &binary) {
+std::optional<mlir::ModuleOp> xilinx::AIE::convertControlPacketBinaryToMLIR(
+    mlir::MLIRContext *ctx, std::vector<uint8_t> &binary, AIEDevice device) {
 
   // parse the binary
   auto maybeOps = parseControlPacket(binary);
@@ -882,13 +881,13 @@ xilinx::AIE::convertControlPacketBinaryToMLIR(mlir::MLIRContext *ctx,
   OpBuilder builder(module.getBodyRegion());
   builder.setInsertionPointToStart(module.getBody());
 
-  // Create aie.device - default to npu1_1col
-  // Note: Control packets don't have device info in the binary format
-  auto device = builder.create<DeviceOp>(loc, AIEDevice::npu1_1col,
-                                         StringAttr::get(builder.getContext()));
-  device.getRegion().emplaceBlock();
-  DeviceOp::ensureTerminator(device.getBodyRegion(), builder, loc);
-  builder.setInsertionPointToStart(device.getBody());
+  // Create aie.device with specified device type
+  auto deviceAttr = AIEDeviceAttr::get(ctx, device);
+  auto deviceOp = builder.create<DeviceOp>(
+      loc, deviceAttr, StringAttr::get(builder.getContext()));
+  deviceOp.getRegion().emplaceBlock();
+  DeviceOp::ensureTerminator(deviceOp.getBodyRegion(), builder, loc);
+  builder.setInsertionPointToStart(deviceOp.getBody());
 
   // Create tiles and set controller_id attributes based on packet info
   // Group operations by (pktType, pktId) to determine which tile they target
@@ -897,13 +896,13 @@ xilinx::AIE::convertControlPacketBinaryToMLIR(mlir::MLIRContext *ctx,
     auto key = std::make_pair(op.pktType, op.pktId);
     if (tileMap.find(key) == tileMap.end()) {
       // Extract col/row from address using target model
-      const AIETargetModel &targetModel = device.getTargetModel();
+      const AIETargetModel &targetModel = deviceOp.getTargetModel();
       uint32_t colInt = (op.address >> targetModel.getColumnShift()) & 0x1f;
       uint32_t rowInt = (op.address >> targetModel.getRowShift()) & 0x1f;
       tileMap[key] = std::make_pair(colInt, rowInt);
 
       // Create tile and set controller_id attribute
-      auto tile = TileOp::getOrCreate(builder, device, colInt, rowInt);
+      auto tile = TileOp::getOrCreate(builder, deviceOp, colInt, rowInt);
       auto packetInfoAttr =
           AIE::PacketInfoAttr::get(builder.getContext(), op.pktType, op.pktId);
       tile->setAttr("controller_id", packetInfoAttr);

python/AIEMLIRModule.cpp

@@ -135,12 +135,12 @@ NB_MODULE(_aie, m) {
 
   m.def(
       "control_packets_binary_to_mlir",
-      [](MlirContext ctx, nb::bytes bytes) {
+      [](MlirContext ctx, nb::bytes bytes, int device) {
         MlirStringRef bin = {static_cast<const char *>(bytes.data()),
                              bytes.size()};
-        return aieTranslateBinaryToControlPackets(ctx, bin);
+        return aieTranslateBinaryToControlPackets(ctx, bin, device);
       },
-      "ctx"_a, "binary"_a);
+      "ctx"_a, "binary"_a, "device"_a = 4); /* npu1 */
 
   m.def(
       "translate_npu_to_binary",

python/compiler/pkt2mlir/README.md

@@ -12,12 +12,14 @@ The tool parses binary control packet files generated by `aie-translate --aie-ct
 ## Usage
 
 ```bash
-pkt2mlir.py -f <binary-file>
+pkt2mlir.py -f <binary-file> [-d <device-type>]
 ```
 
 ### Options
 
 - `-f`, `-file FILE` : Input control packet binary file (required)
+- `-d`, `-device DEVICE` : Target AIE device type (default: npu1)
+  - Supported devices: xcvc1902, xcve2302, xcve2802, npu1, npu1_1col, npu1_2col, npu1_3col, npu2, npu2_1col, npu2_2col, npu2_3col, npu2_4col, npu2_5col, npu2_6col, npu2_7col
 - `-h`, `--help` : Show help message and exit
 
 ## Examples
@@ -65,9 +67,17 @@ Translate back to MLIR:
 pkt2mlir.py -f control_packets.bin
 ```
 
+Or specify a different device type:
+
+```bash
+pkt2mlir.py -f control_packets.bin -d npu1
+```
+
+Output:
+
 ```mlir
 module {
-  aie.device(npu1_1col) {
+  aie.device(npu1) {
     %shim_noc_tile_0_0 = aie.tile(0, 0) {controller_id = #aie.packet_info<pkt_type = 2, pkt_id = 3>}
     
     aiex.runtime_sequence @configure() {

python/compiler/pkt2mlir/main.py

@@ -24,14 +24,24 @@ def main():
         required=True,
         help="Input control packet binary file",
     )
+    parser.add_argument(
+        "-device",
+        "-d",
+        type=str,
+        default="npu1",
+        help="Target AIE device type (default: npu1)",
+    )
 
     args = parser.parse_args()
 
     # Read the data from the file
     data = args.file.read()
 
+    # Get the device enum value
+    device_value = getattr(AIEDevice, args.device)
+
     with Context() as ctx:
-        module = control_packets_binary_to_mlir(ctx, data)
+        module = control_packets_binary_to_mlir(ctx, data, device_value)
 
     print(str(module))