[IR] Add connection field to write_bd and dma_start ops for pairing (#1339)

This enables passing the corresponding BD IDs between paired ops, which
is required for out-of-order DMA mode.

Author: Yu-Zhewen

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td

@@ -1037,12 +1037,14 @@ def AMDAIE_NpuWriteBdOp: AMDAIE_Op<"npu.write_bd"> {
   let summary = "Initialize the buffer descriptor with specified ID";
   let description = [{
     This NPU controller operation to initialize the `bd_id` buffer descriptor on
-    the (`col`, `row`) tile with the provided configurations.
+    the (`col`, `row`) tile with the provided configurations. The connection field
+    is required for the out-of-order DMA mode, to pass to the corresponding BD IDs
+    between paired write_bd/dma_start ops.
 
     Example:
 
     ```mlir
-    amdaie.npu.write_bd {bd_id = 0 : ui32, buffer_length = 0 : ui32,
+    amdaie.npu.write_bd(%0) {bd_id = 0 : ui32, buffer_length = 0 : ui32,
       buffer_offset = 0 : ui32, col = 0 : ui32, enable_packet = false,
       iteration_current = 0 : ui32, iteration_size = 0 : ui32,
       iteration_stride = 0 : ui32, lock_acq_enable = false,
@@ -1054,7 +1056,8 @@ def AMDAIE_NpuWriteBdOp: AMDAIE_Op<"npu.write_bd"> {
     ```
   }];
   let arguments = (
-    ins UI32Attr:$col,
+    ins Optional<Index>:$connection,
+        UI32Attr:$col,
         UI32Attr:$row,
         UI32Attr:$bd_id,
         UI32Attr:$buffer_length,
@@ -1079,7 +1082,7 @@ def AMDAIE_NpuWriteBdOp: AMDAIE_Op<"npu.write_bd"> {
         I32Attr:$lock_acq_val,
         UI32Attr:$lock_acq_id
   );
-  let assemblyFormat = [{ attr-dict }];
+  let assemblyFormat = [{ ( `(` $connection^ `)` ) ? attr-dict }];
 }
 
 def AMDAIE_NpuTctSyncOp: AMDAIE_Op<"npu.tct_sync"> {
@@ -1749,16 +1752,25 @@ def AMDAIE_DMAStartOp: AMDAIE_Op<"dma_start", [
   ]>, Results<(outs I1:$valid)> {
 
   let summary = "An op to start DMA";
+  let description = [{
+    This opertion specifies a list of DMA tasks (BDs) to be exectuted on the given channel.
+    The connection field is required for the out-of-order DMA mode, to pass to the corresponding
+    BD IDs between paired write_bd/dma_start ops.
+  }];
+
   let arguments = (
-    ins Index:$tile,
-        DMAChannelDir:$channel_dir,
-        ConfinedAttr<I8Attr, [IntMinValue<0>]>:$channel_index,
+    ins Index:$channel,
+        Variadic<Index>:$connections,
         // `repeat_count==1` means "do it once".
         DefaultValuedAttr<I8Attr, "1">:$repeat_count
   );
   let regions = (region AnyRegion:$body);
   let assemblyFormat = [{
-    `(` $tile `,` $channel_dir `,` $channel_index (`,` `repeat_count` `=` $repeat_count^)? `)` regions attr-dict
+    `(`
+    $channel
+    ( `,` `{` $connections^ `}` )?
+    `)`
+    regions attr-dict
   }];
 }
 

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/roundtrip.mlir

@@ -468,9 +468,11 @@ func.func @npu_push_to_queue() {
 // -----
 
 // CHECK-LABEL: func.func @npu_write_bd
-// CHECK:       amdaie.npu.write_bd {bd_id = 2 : ui32, buffer_length = 1024 : ui32, buffer_offset = 32 : ui32, col = 1 : ui32, enable_packet = true, iteration_current = 0 : ui32, iteration_size = 0 : ui32, iteration_stride = 0 : ui32, lock_acq_enable = false, lock_acq_id = 0 : ui32, lock_acq_val = 0 : i32, lock_rel_id = 0 : ui32, lock_rel_val = 0 : i32, next_bd = 0 : ui32, out_of_order_id = 0 : ui32, packet_id = 1 : ui32, packet_type = 0 : ui32, paddings_after = array<i32>, paddings_before = array<i32>, row = 0 : ui32, sizes = array<i32: 4, 16, 16>, strides = array<i32: 64, 8, 1>, use_next_bd = false, valid_bd = true}
-func.func @npu_write_bd() {
-  amdaie.npu.write_bd {bd_id = 2 : ui32, buffer_length = 1024 : ui32, buffer_offset = 32 : ui32, col = 1 : ui32, enable_packet = true, iteration_current = 0 : ui32, iteration_size = 0 : ui32, iteration_stride = 0 : ui32, lock_acq_enable = false, lock_acq_id = 0 : ui32, lock_acq_val = 0 : i32, lock_rel_id = 0 : ui32, lock_rel_val = 0 : i32, next_bd = 0 : ui32, out_of_order_id = 0 : ui32, packet_id = 1 : ui32, packet_type = 0 : ui32, paddings_after = array<i32>, paddings_before = array<i32>, row = 0 : ui32, sizes = array<i32: 4, 16, 16>, strides = array<i32: 64, 8, 1>, use_next_bd = false, valid_bd = true}
+// CHECK:         %[[CONNECTION_0:.+]] = amdaie.connection
+// CHECK:         amdaie.npu.write_bd(%[[CONNECTION_0]]) {bd_id = 2 : ui32, buffer_length = 1024 : ui32, buffer_offset = 32 : ui32, col = 1 : ui32, enable_packet = true, iteration_current = 0 : ui32, iteration_size = 0 : ui32, iteration_stride = 0 : ui32, lock_acq_enable = false, lock_acq_id = 0 : ui32, lock_acq_val = 0 : i32, lock_rel_id = 0 : ui32, lock_rel_val = 0 : i32, next_bd = 0 : ui32, out_of_order_id = 0 : ui32, packet_id = 1 : ui32, packet_type = 0 : ui32, paddings_after = array<i32>, paddings_before = array<i32>, row = 0 : ui32, sizes = array<i32: 4, 16, 16>, strides = array<i32: 64, 8, 1>, use_next_bd = false, valid_bd = true}
+func.func @npu_write_bd(%arg0: !amdaie.logicalobjectfifo<memref<32xi32>>, %arg1: !amdaie.logicalobjectfifo<memref<32xi32, 1 : i32>>) {
+  %0 = amdaie.connection(%arg0, %arg1) : (!amdaie.logicalobjectfifo<memref<32xi32>>, !amdaie.logicalobjectfifo<memref<32xi32, 1 : i32>>)
+  amdaie.npu.write_bd(%0) {bd_id = 2 : ui32, buffer_length = 1024 : ui32, buffer_offset = 32 : ui32, col = 1 : ui32, enable_packet = true, iteration_current = 0 : ui32, iteration_size = 0 : ui32, iteration_stride = 0 : ui32, lock_acq_enable = false, lock_acq_id = 0 : ui32, lock_acq_val = 0 : i32, lock_rel_id = 0 : ui32, lock_rel_val = 0 : i32, next_bd = 0 : ui32, out_of_order_id = 0 : ui32, packet_id = 1 : ui32, packet_type = 0 : ui32, paddings_after = array<i32>, paddings_before = array<i32>, row = 0 : ui32, sizes = array<i32: 4, 16, 16>, strides = array<i32: 64, 8, 1>, use_next_bd = false, valid_bd = true}
   return
 }
 
@@ -678,7 +680,9 @@ func.func @from_memref_unknown_row_column(%arg0 : memref<8xi32>, %t0 : index) {
 //       CHECK: %[[LOCK_0:.*]] = amdaie.lock
 //       CHECK: %[[LOCK_1:.*]] = amdaie.lock
 //       CHECK: %[[BUFFER:.*]] = amdaie.buffer
-//       CHECK: amdaie.dma_start(%[[TILE]], S2MM, 1) {
+//       CHECK: %[[CHANNEL:.*]] = amdaie.channel
+//       CHECK: %[[CONNECTION:.*]] = amdaie.connection
+//       CHECK: amdaie.dma_start(%[[CHANNEL]], {%[[CONNECTION]]}) {
 //       CHECK:   amdaie.use_lock(%[[LOCK_0]], AcquireGreaterOrEqual(1))
 //       CHECK:   amdaie.dma_bd(%[[BUFFER]] : memref<1024xi32, 2 : i32>)
 //  CHECK-SAME:     {dimensions = #amdaie<bd_dim_layout_array[<size = 32, stride = 8>, <size = 4, stride = 256>, <size = 8, stride = 1>]>, len = 1024 : i32}
@@ -687,14 +691,16 @@ func.func @from_memref_unknown_row_column(%arg0 : memref<8xi32>, %t0 : index) {
 //       CHECK:  ^bb1:
 //       CHECK:   amdaie.end
 //       CHECK: }
-func.func @dma_start_block_ops() {
+func.func @dma_start_block_ops(%arg0: !amdaie.logicalobjectfifo<memref<1024xi32, 1 : i32>>, %arg1: !amdaie.logicalobjectfifo<memref<1024xi32, 2 : i32>>) {
   %c0 = arith.constant 0 : index
   %c2 = arith.constant 2 : index
   %tile_0_2 = amdaie.tile(%c0, %c2)
   %lock = amdaie.lock(%tile_0_2(0), 1)
   %lock_0 = amdaie.lock(%tile_0_2(1), 0)
   %buffer = amdaie.buffer(%tile_0_2) : memref<1024xi32, 2 : i32>
-  %0 = amdaie.dma_start(%tile_0_2, S2MM, 1) {
+  %channel = amdaie.channel(%tile_0_2, 0, port_type = DMA, direction = MM2S)
+  %connection = amdaie.connection(%arg0, %arg1 {%channel}) : (!amdaie.logicalobjectfifo<memref<1024xi32, 1 : i32>>, !amdaie.logicalobjectfifo<memref<1024xi32, 2 : i32>>)
+  amdaie.dma_start(%channel, {%connection}) {
     amdaie.use_lock(%lock, AcquireGreaterOrEqual(1))
     amdaie.dma_bd(%buffer : memref<1024xi32, 2 : i32>) {dimensions = #amdaie<bd_dim_layout_array[<size = 32, stride = 8>, <size = 4, stride = 256>, <size = 8, stride = 1>]>, len = 1024 : i32}
     amdaie.use_lock(%lock_0, Release(1))

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignBDIDs.cpp

@@ -52,7 +52,8 @@ LogicalResult assignBdIds(Operation *deviceOp) {
   deviceOp->walk(
       [&](AMDAIE::DMAStartOp dmaStartOp) { memOps.push_back(dmaStartOp); });
   for (AMDAIE::DMAStartOp dmaStartOp : memOps) {
-    auto tile = dmaStartOp.getTile().getDefiningOp<AMDAIE::TileOp>();
+    auto channelOp = dmaStartOp.getChannel().getDefiningOp<AMDAIE::ChannelOp>();
+    AMDAIE::TileOp tile = channelOp.getTileOp();
     std::optional<int64_t> col = getConstantIntValue(tile.getCol());
     std::optional<int64_t> row = getConstantIntValue(tile.getRow());
     if (!col || !row) {
@@ -69,8 +70,8 @@ LogicalResult assignBdIds(Operation *deviceOp) {
 
     DenseMap<Block *, int> blockChannelMap;
     // Associate with each block the channel index specified by the
-    // dma_start
-    int chNum = dmaStartOp.getChannelIndex();
+    // dma_start.
+    int chNum = channelOp.getValue();
     dmaStartOp->walk<WalkOrder::PreOrder>([&](AMDAIE::DMABDOp bd) {
       if (bd.getBdId().has_value()) {
         assert(gen.isBdIdAssigned(bd.getBdId().value()) &&

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEControlCodeLowering.cpp

@@ -35,11 +35,11 @@ struct HalfDmaCpyNdToNpuConverter final
   /// is specific to Shim BDs for now.
   FailureOr<AMDAIE::NpuPushToQueueOp> insertWriteBdOps(
       AMDAIE::NpuHalfDmaCpyNdOp op, ConversionPatternRewriter &rewriter,
-      AMDAIE::AMDAIETileType tileType, AMDAIE::BdIdOp bdIdOp,
-      AMDAIE::ChannelOp channelOp, int64_t bufferLength, int64_t bufferOffset,
-      int32_t enablePacket, int32_t packetId, int32_t packetType,
-      SmallVector<OpFoldResult> sizes, SmallVector<OpFoldResult> strides,
-      bool loweringCtrlpktDma) const {
+      AMDAIE::AMDAIETileType tileType, AMDAIE::ConnectionOp connectionOp,
+      AMDAIE::BdIdOp bdIdOp, AMDAIE::ChannelOp channelOp, int64_t bufferLength,
+      int64_t bufferOffset, int32_t enablePacket, int32_t packetId,
+      int32_t packetType, SmallVector<OpFoldResult> sizes,
+      SmallVector<OpFoldResult> strides, bool loweringCtrlpktDma) const {
     FailureOr<uint8_t> maybeNumIntraAddrDim = deviceModel.getDmaProp<uint8_t>(
         tileType, AMDAIE::AMDAIEDmaProp::NumAddrDim);
     if (failed(maybeNumIntraAddrDim)) return failure();
@@ -190,11 +190,11 @@ struct HalfDmaCpyNdToNpuConverter final
     // Offset set to zero for shim as the offset is embedded in the address
     // patch.
     rewriter.create<AMDAIE::NpuWriteBdOp>(
-        op.getLoc(), col, row, bdId, innerBufferLength, 0, staticSizes,
-        staticStrides, paddingsBefore, paddingsAfter, iterationCurrent,
-        iterationSize, iterationStride, enablePacket, packetId, packetType,
-        outOfOrderId, useNextBd, nextBd, validBd, lockAcqEnable, lockRelVal,
-        lockRelId, lockAcqVal, lockAcqId);
+        op.getLoc(), connectionOp, col, row, bdId, innerBufferLength, 0,
+        staticSizes, staticStrides, paddingsBefore, paddingsAfter,
+        iterationCurrent, iterationSize, iterationStride, enablePacket,
+        packetId, packetType, outOfOrderId, useNextBd, nextBd, validBd,
+        lockAcqEnable, lockRelVal, lockRelId, lockAcqVal, lockAcqId);
     rewriter.create<AMDAIE::NpuAddressPatchOp>(op.getLoc(), col, bdId, argIdx,
                                                bufferOffsetInBytes);
     SmallVector<Type> resultTypes = {
@@ -267,10 +267,10 @@ struct HalfDmaCpyNdToNpuConverter final
     SmallVector<OpFoldResult> sizes = op.getMixedSizes();
     SmallVector<OpFoldResult> strides = op.getMixedStrides();
     FailureOr<AMDAIE::NpuPushToQueueOp> npuPushToQueueOp = insertWriteBdOps(
-        op, rewriter, AMDAIE::AMDAIETileType::SHIMNOC, maybeBdIdOp.value(),
-        maybeChannelOp.value(), maybeSize.value(), maybeOffset.value(),
-        enablePacket, packetId, packetType, sizes, strides,
-        *maybeIsControlFlow);
+        op, rewriter, AMDAIE::AMDAIETileType::SHIMNOC,
+        maybeConnectionOp.value(), maybeBdIdOp.value(), maybeChannelOp.value(),
+        maybeSize.value(), maybeOffset.value(), enablePacket, packetId,
+        packetType, sizes, strides, *maybeIsControlFlow);
     if (failed(npuPushToQueueOp)) return failure();
     rewriter.replaceOp(op, *npuPushToQueueOp);
 
@@ -411,17 +411,17 @@ BDDimLayoutAndLength convertSizeStrideToBDDimLayoutArrayAttr(
 ///       amdaie.end
 ///  }
 LogicalResult createDMABlocks(
-    IRRewriter &rewriter, Operation *tileOp, AMDAIE::DMAChannelDir channelDir,
-    int channelIndex, ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides,
-    size_t acqNum, size_t relNum, int64_t offset,
+    IRRewriter &rewriter, AMDAIE::ChannelOp channelOp,
+    AMDAIE::ConnectionOp connectionOp, ArrayRef<int64_t> sizes,
+    ArrayRef<int64_t> strides, size_t acqNum, size_t relNum, int64_t offset,
     const SmallVector<AMDAIE::BufferOp> &bufferOps,
     const std::pair<AMDAIE::LockOp, AMDAIE::LockOp> &locks,
     std::optional<uint8_t> pktId) {
   OpBuilder::InsertionGuard g(rewriter);
 
   // Create DMA start on channel.
   auto dmaStartOp = rewriter.create<AMDAIE::DMAStartOp>(
-      rewriter.getUnknownLoc(), tileOp->getResult(0), channelDir, channelIndex,
+      rewriter.getUnknownLoc(), channelOp, ValueRange{connectionOp},
       /*repeatCount=*/1);
   rewriter.setInsertionPointToStart(&dmaStartOp.getRegion().emplaceBlock());
   rewriter.create<AMDAIE::EndOp>(rewriter.getUnknownLoc());
@@ -583,10 +583,10 @@ static LogicalResult processConnectionSourceForDmaStart(
             sourceMemSpace, [&]() { return connectionOp->emitOpError(); }))) {
       return failure();
     };
-    if (failed(createDMABlocks(rewriter, tileOp, AMDAIE::DMAChannelDir::MM2S,
-                               channel.getValue(), canonicalizedSizes,
-                               canonicalizedStrides, acqNum, acqNum,
-                               sourceOffset, buffers, lockPair, packetId))) {
+    if (failed(createDMABlocks(rewriter, channel, connectionOp,
+                               canonicalizedSizes, canonicalizedStrides, acqNum,
+                               acqNum, sourceOffset, buffers, lockPair,
+                               packetId))) {
       return sourceObjFifo.emitOpError() << "could not create DMA operations";
     }
   }
@@ -666,10 +666,10 @@ static LogicalResult processConnectionTargetForDmaStart(
             targetMemSpace, [&]() { return connectionOp->emitOpError(); }))) {
       return failure();
     };
-    if (failed(createDMABlocks(rewriter, tileOp, AMDAIE::DMAChannelDir::S2MM,
-                               channel.getValue(), canonicalizedSizes,
-                               canonicalizedStrides, acqNum, acqNum,
-                               targetOffset, buffers, lockPair, packetId))) {
+    if (failed(createDMABlocks(rewriter, channel, connectionOp,
+                               canonicalizedSizes, canonicalizedStrides, acqNum,
+                               acqNum, targetOffset, buffers, lockPair,
+                               packetId))) {
       return targetObjFifo.emitOpError() << "could not create DMA operations";
     }
   }

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Utils/AMDAIETransactionBuilder.cpp

@@ -84,15 +84,20 @@ LogicalResult TransactionBuilder::appendLockOp(AMDAIE::LockOp lockOp) {
 
 LogicalResult TransactionBuilder::appendDmaStartOp(
     AMDAIE::DMAStartOp dmaStartOp) {
-  // Configure DMA Locks.
-  auto tile = dmaStartOp.getTile().getDefiningOp<AMDAIE::TileOp>();
+  // Get tile location.
+  auto channelOp = dmaStartOp.getChannel().getDefiningOp<AMDAIE::ChannelOp>();
+  AMDAIE::TileOp tile = channelOp.getTileOp();
   std::optional<int64_t> maybeCol = getConstantIntValue(tile.getCol());
   std::optional<int64_t> maybeRow = getConstantIntValue(tile.getRow());
   if (!maybeCol || !maybeRow) {
     return tile->emitOpError()
            << "expected column and row integer value/constant";
   }
   XAie_LocType tileLoc = XAie_TileLoc(*maybeCol, *maybeRow);
+  // Get channel op.
+  int chNum = channelOp.getValue();
+  auto channelDir = static_cast<DMAChannelDir>(channelOp.getDirection());
+  // Initialize the DMA descriptor.
   FailureOr<XAie_DmaDesc> dmaDesc = initDMADesc(deviceModel, tileLoc);
   if (failed(dmaDesc)) return failure();
 
@@ -172,8 +177,6 @@ LogicalResult TransactionBuilder::appendDmaStartOp(
   // Configure push to BD queue.
   // TODO: Generalize it as this is currently hardcoded to only shim side for
   // now.
-  int chNum = dmaStartOp.getChannelIndex();
-  auto channelDir = static_cast<DMAChannelDir>(dmaStartOp.getChannelDir());
   bool issueToken = tileLoc.Row == 0 && channelDir == DMAChannelDir::MM2S;
   bool setChannelEnable = true;
   if (failed(configurePushToBdQueue(

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_bd_ids.mlir

@@ -6,7 +6,8 @@
 // CHECK:           %[[IN:.*]] = amdaie.buffer(%[[TILE_2_1]])
 // CHECK:           %[[LOCK_2_1:.*]] = amdaie.lock(%[[TILE_2_1]](1), 0)
 // CHECK:           %[[LOCK_2_1_0:.*]] = amdaie.lock(%[[TILE_2_1]](0), 1)
-// CHECK:           amdaie.dma_start(%[[TILE_2_1]], MM2S, 0) {
+// CHECK:           %[[CHANNEL:.*]] = amdaie.channel(%[[TILE_2_1]], 0, port_type = DMA, direction = MM2S)
+// CHECK:           amdaie.dma_start(%[[CHANNEL]]) {
 // CHECK:             amdaie.use_lock(%[[LOCK_2_1]], AcquireGreaterOrEqual
 // CHECK:             amdaie.dma_bd(%[[IN]] : memref<16xi32>) {bd_id = 0 : i32
 // CHECK:             amdaie.next_bd ^bb1
@@ -34,7 +35,8 @@ module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb}
     %buf01_0 = amdaie.buffer(%t01) { address = 8192 : i32, sym_name = "in" } : memref<16xi32>
     %l01_0 = amdaie.lock(%t01(1), 0)
     %l01_1 = amdaie.lock(%t01(0), 1)
-    %dstDma = amdaie.dma_start(%t01, MM2S, 0) {
+    %channel = amdaie.channel(%t01, 0, port_type = DMA, direction = MM2S)
+    %dstDma = amdaie.dma_start(%channel) {
         amdaie.use_lock(%l01_0, AcquireGreaterOrEqual(1))
         amdaie.dma_bd(%buf01_0 : memref<16xi32>) {len = 16 : i32, dimensions = #amdaie<bd_dim_layout_array[<size = 2, stride = 1>, <size = 3, stride = 2>, <size = 2, stride = 4>, <size = 1, stride = 1>]>}
         amdaie.next_bd ^bb1