[AIEX] Add symbol-based IpuDmaWaitOp (#1220)

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

@@ -478,7 +478,8 @@ def AIE_IpuDmaMemcpyNdOp: AIEX_Op<"ipu.dma_memcpy_nd", [
     symbol and specifies a descriptor `id` to be used, which will become the `bd_id` to be used
     when lowered further. The `issue_token` attribute specifies whether the execution of this
     operation should issue a token which can be received and read for synchronization purposes.
-    This `issue_token` attribute is set to `false` by default for `MM2S` for backward compatibility and **is always set to true for** `S2MM` channels.
+    This `issue_token` attribute is set to `false` by default for `MM2S` for backward compatibility
+    and **is always set to true for** `S2MM` channels.
   }];
 
   let arguments = (
@@ -518,6 +519,35 @@ def AIE_IpuDmaMemcpyNdOp: AIEX_Op<"ipu.dma_memcpy_nd", [
   let hasVerifier = 1;
 }
 
+def AIE_IpuDmaWaitOp: AIEX_Op<"ipu.dma_wait", []> {
+  let summary = "Blocking operation to wait for a DMA to complete execution.";
+  let description = [{
+    The IpuDmaWaitOp blocks until the DMA referenced through `symbol` completes execution
+    and issues a task-complete-token.
+
+    Example:
+    ```mlir
+      ...
+      aie.objectfifo @out0(%tile_0_1, {%tile_0_0}, 4 : i32) : !aie.objectfifo<memref<32x32xi32>>
+      ...
+      aiex.ipu.dma_memcpy_nd(0, 0, %arg2[1, 1, 0, 0][1, 1, 32, 32][1, 1, 64]) {id = 0 : i64, issue_token = true, metadata = @out0} : memref<32x64xi32>
+      ...
+      aiex.ipu.dma_wait { symbol = @out0 }
+    ```
+    Here, we have an objectfifo with symbol name `out0`, which is then referenced in the
+    `ipu.dma_memcpy_nd` operation as the target for the respective DMA operation. Afterwards,
+    an `ipu.dma_wait` operation references the same symbol to block until the respective DMA
+    has executed all of its tasks.
+  }];
+  let arguments = (
+    ins FlatSymbolRefAttr:$symbol
+  );
+  let assemblyFormat = [{
+    attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
 // Write RTP
 def AIE_IpuWriteRTPOp: AIEX_Op<"ipu.rtp_write", []> {
   let summary = "rtp write operator";

lib/Dialect/AIEX/IR/AIEXDialect.cpp

@@ -105,6 +105,15 @@ LogicalResult AIEX::IpuDmaMemcpyNdOp::verify() {
   return success();
 }
 
+LogicalResult AIEX::IpuDmaWaitOp::verify() {
+  AIE::DeviceOp dev = (*this)->getParentOfType<AIE::DeviceOp>();
+  if (!dev)
+    return emitOpError("couldn't find parent of type DeviceOp");
+  if (!dev.lookupSymbol(getSymbol()))
+    return emitOpError("couldn't find symbol in parent device");
+  return success();
+}
+
 LogicalResult AIEX::IpuShimTilePushQueueOp::verify() {
   const auto &targetModel = AIE::getTargetModel(*this);
   auto numBds = targetModel.getNumBDs(0, 0); // assume shim

lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp

@@ -335,6 +335,38 @@ struct DmaToIpuPattern : OpConversionPattern<IpuDmaMemcpyNdOp> {
   }
 };
 
+/// Convert IpuDmaWaitOp into IpuSyncOp by retrieving the necessary
+/// information from the ShimDMAAllocationOp referenced through the
+/// symbol argument of this op.
+struct DmaWaitToIpuPattern : OpConversionPattern<IpuDmaWaitOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  DmaWaitToIpuPattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(context, benefit) {}
+
+  LogicalResult
+  matchAndRewrite(IpuDmaWaitOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    AIE::DeviceOp dev = op->getParentOfType<AIE::DeviceOp>();
+    std::optional<AIE::ShimDMAAllocationOp> shimDmaAllocOp =
+        getAllocOpForSymbol(dev, op.getSymbol());
+    if (!shimDmaAllocOp) {
+      op.emitOpError("couldn't find shim_dma_allocation op");
+      return failure();
+    }
+    AIE::DMAChannelDir channelDir = shimDmaAllocOp->getChannelDir();
+    int channel = shimDmaAllocOp->getChannelIndex();
+    int direction = (int)(channelDir == AIE::DMAChannelDir::MM2S);
+    int column = shimDmaAllocOp->getCol();
+
+    // Create with `column_num == 1` and `row_num == 1` to check for a single
+    // column and row. Row is always 0 for shim tiles.
+    (void)rewriter.replaceOpWithNewOp<IpuSyncOp>(op, column, 0, direction,
+                                                 channel, 1, 1);
+    return success();
+  }
+};
+
 struct AIEDmaToIpuPass : AIEDmaToIpuBase<AIEDmaToIpuPass> {
   void runOnOperation() override {
 
@@ -346,10 +378,12 @@ struct AIEDmaToIpuPass : AIEDmaToIpuBase<AIEDmaToIpuPass> {
     target.addLegalOp<AIE::ShimDMAAllocationOp>();
     target.addIllegalOp<IpuWriteRTPOp>();
     target.addIllegalOp<IpuDmaMemcpyNdOp>();
+    target.addIllegalOp<IpuDmaWaitOp>();
     target.addIllegalOp<IpuShimTilePushQueueOp>();
 
     RewritePatternSet patterns(&getContext());
     patterns.insert<DmaToIpuPattern>(&getContext());
+    patterns.insert<DmaWaitToIpuPattern>(&getContext());
     patterns.insert<PushToIpuPattern>(&getContext());
     patterns.insert<RtpToIpuPattern>(&getContext());
 

test/Conversion/DmaToIpu/dma_to_ipu.mlir

@@ -8,29 +8,78 @@
 //
 //===----------------------------------------------------------------------===//
 
-// RUN: aie-opt -aie-dma-to-ipu %s | FileCheck %s
+// RUN: aie-opt --split-input-file -aie-dma-to-ipu %s | FileCheck %s
 
 // TODO - more
-// CHECK-LABEL: test0
+// CHECK-LABEL: dma_memcpy_nd_0
 // CHECK: aiex.ipu.writebd_shimtile
 // CHECK-SAME: ddr_id = 0 : i32
 // CHECK-SAME: valid_bd = 1 : i32
 // CHECK: aiex.ipu.writebd_shimtile
 // CHECK-SAME: ddr_id = 1 : i32
 module  {
-aie.device(ipu) {
-  memref.global "public" @toMem : memref<16xi32>
-  memref.global "public" @fromMem : memref<16xi32>
-  func.func @test0(%arg0: memref<16xi32>, %arg1: memref<16xi32>) {
-    %c16_i64 = arith.constant 16 : i64
-    %c64_i64 = arith.constant 64 : i64
-    %c0_i64 = arith.constant 0 : i64
-    %c1_i64 = arith.constant 1 : i64
-    aiex.ipu.dma_memcpy_nd(0, 0, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64][%c1_i64, %c1_i64, %c16_i64, %c16_i64][%c0_i64, %c0_i64, %c64_i64]) { metadata = @toMem, id = 1 : i64 } : memref<16xi32>
-    aiex.ipu.dma_memcpy_nd(0, 1, %arg1[%c0_i64, %c0_i64, %c0_i64, %c16_i64][%c1_i64, %c1_i64, %c16_i64, %c16_i64][%c0_i64, %c0_i64, %c64_i64]) { metadata = @fromMem, id = 0 : i64 } : memref<16xi32>
-    return
+  aie.device(ipu) {
+    memref.global "public" @toMem : memref<16xi32>
+    memref.global "public" @fromMem : memref<16xi32>
+    func.func @dma_memcpy_nd_0(%arg0: memref<16xi32>, %arg1: memref<16xi32>) {
+      aiex.ipu.dma_memcpy_nd(0, 0, %arg0[0, 0, 0, 0][1, 1, 16, 16][0, 0, 64]) { metadata = @toMem, id = 1 : i64 } : memref<16xi32>
+      aiex.ipu.dma_memcpy_nd(0, 1, %arg1[0, 0, 0, 16][1, 1, 16, 16][0, 0, 64]) { metadata = @fromMem, id = 0 : i64 } : memref<16xi32>
+      return
+    }
+    aie.shim_dma_allocation @fromMem (MM2S, 0, 0)
+    aie.shim_dma_allocation @toMem (S2MM, 0, 0)
   }
-  aie.shim_dma_allocation @fromMem (MM2S, 0, 0)
-  aie.shim_dma_allocation @toMem (S2MM, 0, 0)
 }
+
+// -----
+
+// CHECK-LABEL: dma_wait_s2mm
+// CHECK: aiex.ipu.writebd_shimtile
+// CHECK-SAME: ddr_id = 0 : i32
+// CHECK-SAME: valid_bd = 1 : i32
+// CHECK: aiex.ipu.write32
+// CHECK: aiex.ipu.sync 
+// CHECK-SAME: channel = 0 : i32
+// CHECK-SAME: column = 0 : i32
+// CHECK-SAME: column_num = 1 : i32
+// CHECK-SAME: direction = 0 : i32
+// CHECK-SAME: row = 0 : i32
+// CHECK-SAME: row_num = 1 : i32
+module  {
+  aie.device(ipu) {
+    memref.global "public" @toMem : memref<16xi32>
+    func.func @dma_wait_s2mm(%arg0: memref<16xi32>, %arg1: memref<16xi32>) {
+      aiex.ipu.dma_memcpy_nd(0, 0, %arg0[0, 0, 0, 0][1, 1, 16, 16][0, 0, 64]) { metadata = @toMem, id = 1 : i64 } : memref<16xi32>
+      aiex.ipu.dma_wait {symbol = @toMem}
+      return
+    }
+    aie.shim_dma_allocation @toMem (S2MM, 0, 0)
+  }
+}
+
+// -----
+
+// CHECK-LABEL: dma_wait_mm2s
+// CHECK: aiex.ipu.writebd_shimtile
+// CHECK-SAME: ddr_id = 0 : i32
+// CHECK-SAME: valid_bd = 1 : i32
+// CHECK: aiex.ipu.write32
+// CHECK: aiex.ipu.sync 
+// CHECK-SAME: channel = 1 : i32
+// CHECK-SAME: column = 1 : i32
+// CHECK-SAME: column_num = 1 : i32
+// CHECK-SAME: direction = 1 : i32
+// CHECK-SAME: row = 0 : i32
+// CHECK-SAME: row_num = 1 : i32
+module  {
+  aie.device(ipu) {
+    memref.global "public" @toMem : memref<16xi32>
+    func.func @dma_wait_mm2s(%arg0: memref<16xi32>, %arg1: memref<16xi32>) {
+      aiex.ipu.dma_memcpy_nd(0, 0, %arg0[0, 0, 0, 0][1, 1, 16, 16][0, 0, 64]) { metadata = @toMem, id = 1 : i64 } : memref<16xi32>
+      aiex.ipu.dma_wait {symbol = @toMem}
+      return
+    }
+    aie.shim_dma_allocation @toMem (MM2S, 1, 1)
+  }
 }
+

test/Conversion/DmaToIpu/dma_to_ipu_invalid.mlir

@@ -0,0 +1,23 @@
+//===- dma_to_ipu_invalid.mlir ---------------------------------*- MLIR -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// (c) Copyright 2024 Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt --aie-dma-to-ipu --verify-diagnostics %s
+
+module  {
+  aie.device(ipu) {
+    memref.global "public" @toMem : memref<16xi32>
+    func.func @sequence() {
+      // expected-error@+2 {{failed to legalize operation 'aiex.ipu.dma_wait' that was explicitly marked illegal}}
+      // expected-error@+1 {{couldn't find shim_dma_allocation op}}
+      aiex.ipu.dma_wait {symbol = @toMem}
+      return
+    }
+  }
+}

test/dialect/AIEX/invalid.mlir

@@ -0,0 +1,27 @@
+//===- invalid.mlir --------------------------------------------*- MLIR -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (C) 2024, Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt --split-input-file --verify-diagnostics %s
+
+func.func @ipu_dma_wait_no_device() {
+  // expected-error@+1 {{'aiex.ipu.dma_wait' op couldn't find parent of type DeviceOp}}
+  aiex.ipu.dma_wait {symbol = @out0}
+  return
+}
+
+// -----
+
+aie.device(ipu) {
+  func.func @ipu_dma_wait_no_symbol() {
+    // expected-error@+1 {{'aiex.ipu.dma_wait' op couldn't find symbol in parent device}}
+    aiex.ipu.dma_wait {symbol = @out0}
+    return
+  }
+}

test/dialect/AIEX/roundtrip.mlir

@@ -0,0 +1,22 @@
+//===- roundtrip.mlir ------------------------------------------*- MLIR -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (C) 2024, Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func.func @ipu_dma_wait
+// CHECK: aiex.ipu.dma_wait {symbol = @out0}
+
+aie.device(ipu) {
+  memref.global "public" @out0 : memref<16xi32>
+  func.func @ipu_dma_wait() {
+    aiex.ipu.dma_wait {symbol = @out0}
+    return
+  }
+}