Enable packet attribute in aie.dma_bd op (Xilinx#1674)

fifield · Aug 9, 2024 · 058ba76 · 058ba76
1 parent 5019cee
commit 058ba76
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 7 deletions.
diff --git a/include/aie/Dialect/AIE/IR/AIEOps.td b/include/aie/Dialect/AIE/IR/AIEOps.td
@@ -866,6 +866,7 @@ def AIE_DMABDOp: AIE_Op<"dma_bd", []> {
         OptionalAttr<BDPadLayoutArrayAttr>:$pad_dimensions,
         DefaultValuedOptionalAttr<AIEI32Attr, "0">:$pad_value,
         OptionalAttr<AIEI32Attr>:$bd_id,
+        OptionalAttr<PacketInfoAttr>:$packet,
         // should never be assigned by user...
         OptionalAttr<AIEI32Attr>:$next_bd_id
   );
@@ -912,6 +913,34 @@ def AIE_DMABDOp: AIE_Op<"dma_bd", []> {
       $_state.addAttribute("len", $_builder.getI32IntegerAttr(len));
       $_state.addAttribute("dimensions", dims);
       $_state.addAttribute("pad_dimensions", paddims);
+    }]>,
+    OpBuilder<(ins "mlir::Value":$buffer, "int":$offset, "int":$len, "PacketInfoAttr":$pkt), [{
+      $_state.addOperands(buffer);
+      $_state.addAttribute("offset", $_builder.getI32IntegerAttr(offset));
+      $_state.addAttribute("len", $_builder.getI32IntegerAttr(len));
+      $_state.addAttribute("packet", pkt);
+    }]>,
+    OpBuilder<(ins "mlir::Value":$buffer, "int":$offset, "int":$len, "BDDimLayoutArrayAttr":$dims, "PacketInfoAttr":$pkt), [{
+      $_state.addOperands(buffer);
+      $_state.addAttribute("offset", $_builder.getI32IntegerAttr(offset));
+      $_state.addAttribute("len", $_builder.getI32IntegerAttr(len));
+      $_state.addAttribute("dimensions", dims);
+      $_state.addAttribute("packet", pkt);
+    }]>,
+    OpBuilder<(ins "mlir::Value":$buffer, "int":$offset, "int":$len, "BDPadLayoutArrayAttr":$paddims, "PacketInfoAttr":$pkt), [{
+      $_state.addOperands(buffer);
+      $_state.addAttribute("offset", $_builder.getI32IntegerAttr(offset));
+      $_state.addAttribute("len", $_builder.getI32IntegerAttr(len));
+      $_state.addAttribute("pad_dimensions", paddims);
+      $_state.addAttribute("packet", pkt);
+    }]>,
+    OpBuilder<(ins "mlir::Value":$buffer, "int":$offset, "int":$len, "BDDimLayoutArrayAttr":$dims, "BDPadLayoutArrayAttr":$paddims, "PacketInfoAttr":$pkt), [{
+      $_state.addOperands(buffer);
+      $_state.addAttribute("offset", $_builder.getI32IntegerAttr(offset));
+      $_state.addAttribute("len", $_builder.getI32IntegerAttr(len));
+      $_state.addAttribute("dimensions", dims);
+      $_state.addAttribute("pad_dimensions", paddims);
+      $_state.addAttribute("packet", pkt);
     }]>
   ];
 }

diff --git a/lib/Dialect/AIE/IR/AIEDialect.cpp b/lib/Dialect/AIE/IR/AIEDialect.cpp
@@ -1830,6 +1830,12 @@ LogicalResult DMABDOp::verify() {
                << offsetInBytes << " (bytes)";
     }
   }
+  if (auto packetInfo = getPacket()) {
+    if (packetInfo->getPktType() > 7)
+      return emitOpError("Packet type field can only hold 3 bits.");
+    if (packetInfo->getPktId() > 31)
+      return emitOpError("Packet ID field can only hold 5 bits.");
+  }
 
   if (!getLen() && !getBuffer().getType().hasStaticShape())
     return emitOpError() << "buffer with dynamic shape requires static length.";

diff --git a/lib/Targets/AIETargetCDODirect.cpp b/lib/Targets/AIETargetCDODirect.cpp
@@ -248,6 +248,8 @@ LogicalResult configureBdInBlock(XAie_DevInst &devInst, XAie_DmaDesc &dmaTileBd,
                                  std::optional<int> nextBdId) {
   std::optional<int> packetType;
   std::optional<int> packetID;
+
+  // Below should go
   auto maybePacketOps = block.getOps<DMABDPACKETOp>();
   if (!maybePacketOps.empty()) {
     assert(llvm::range_size(maybePacketOps) == 1 &&
@@ -363,6 +365,11 @@ LogicalResult configureBdInBlock(XAie_DevInst &devInst, XAie_DmaDesc &dmaTileBd,
                             nextBdId.value(), enableNextBd);
   }
 
+  if (auto packetInfo = bdOp.getPacket()) {
+    packetType = packetInfo->getPktType();
+    packetID = packetInfo->getPktId();
+  }
+
   if (packetID) {
     if (!packetType)
       bdOp.emitError("must have packetType with packetID");

diff --git a/test/dialect/AIE/bad_dma_op.mlir b/test/dialect/AIE/bad_dma_op.mlir
@@ -8,7 +8,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-// RUN: not aie-opt %s 2>&1 | FileCheck %s
+// RUN: not aie-opt -split-input-file %s 2>&1 | FileCheck %s
 
 // CHECK: error: 'aie.dma' op DMAOp can only appear in single block region
 module {
@@ -32,3 +32,21 @@ module {
 
   }
 }
+
+// -----
+
+// CHECK: error: 'aie.dma_bd' op Packet ID field can only hold 5 bits.
+module {
+  aie.device(npu1_4col) {
+    %tile14 = aie.tile(1, 4)
+    %buf14 = aie.buffer(%tile14) { sym_name = "buf14" } : memref<128xi32>
+    %mem14 = aie.mem(%tile14) {
+      %srcDma = aie.dma_start("MM2S", 0, ^bd0, ^end)
+      ^bd0:
+        aie.dma_bd(%buf14 : memref<128xi32>, 0, 128, [<size = 1, stride = 128>]) {packet = #aie.packet_info<pkt_type = 7, pkt_id = 33>}
+        aie.next_bd ^end
+      ^end: 
+        aie.end
+    }
+  }
+}
diff --git a/test/npu-xrt/packet_flow/aie.mlir b/test/npu-xrt/packet_flow/aie.mlir
@@ -95,14 +95,12 @@ module {
       }]
       %1 = aie.dma(MM2S, 0) [{
         aie.use_lock(%objFifo_in0_cons_cons_lock, AcquireGreaterEqual, 1)
-        aie.dma_bd_packet(0, 0)
-        aie.dma_bd(%objFifo_in0_cons_buff_0 : memref<64x64xi8>)
+        aie.dma_bd(%objFifo_in0_cons_buff_0 : memref<64x64xi8>) {packet = #aie.packet_info<pkt_type = 0, pkt_id = 0>}
         aie.use_lock(%objFifo_in0_cons_prod_lock, Release, 1)
       }]
       %2 = aie.dma(MM2S, 1) [{
         aie.use_lock(%objFifo_out0_cons_lock, AcquireGreaterEqual, 1)
-        aie.dma_bd_packet(0, 2)
-        aie.dma_bd(%objFifo_out0_buff_0 : memref<64x64xi8>)
+        aie.dma_bd(%objFifo_out0_buff_0 : memref<64x64xi8>) {packet = #aie.packet_info<pkt_type = 0, pkt_id = 2>}
         aie.use_lock(%objFifo_out0_prod_lock, Release, 1)
       }]
       %3 = aie.dma(S2MM, 1) [{
@@ -123,8 +121,7 @@ module {
       }]
       %1 = aie.dma(MM2S, 0) [{
         aie.use_lock(%objFifo_out1_cons_lock, AcquireGreaterEqual, 1)
-        aie.dma_bd_packet(0, 1)
-        aie.dma_bd(%objFifo_out1_buff_0 : memref<64x64xi8>)
+        aie.dma_bd(%objFifo_out1_buff_0 : memref<64x64xi8>) {packet = #aie.packet_info<pkt_type = 0, pkt_id = 1>}
         aie.use_lock(%objFifo_out1_prod_lock, Release, 1)
       }]
       aie.end