Add issue_token attribute to IpuDmaMemcpyNdOp (#1209)

Co-authored-by: Maksim Levental <[email protected]>

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

@@ -470,7 +470,15 @@ def AIE_IpuDmaMemcpyNdOp: AIEX_Op<"ipu.dma_memcpy_nd", [
   let summary = "half dma operator";
 
   let description = [{
-    nd half dma operator
+    An nd half dma operator.
+
+    Programs a DMA on coordinates (`x`, `y`) to access a memory `memref` with an access
+    pattern specified by `offsets`, `sizes` and `strides` or `static_offsets`, `static_sizes`
+    and `static_strides`. The operator references the target channel through the `metadata`
+    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.
   }];
 
   let arguments = (
@@ -485,7 +493,8 @@ def AIE_IpuDmaMemcpyNdOp: AIEX_Op<"ipu.dma_memcpy_nd", [
         ConfinedAttr<DenseI64ArrayAttr, [DenseArrayCount<4>]>:$static_sizes,
         ConfinedAttr<DenseI64ArrayAttr, [DenseArrayCount<3>]>:$static_strides,
         FlatSymbolRefAttr:$metadata,
-        I64Attr:$id
+        I64Attr:$id,
+        DefaultValuedOptionalAttr<BoolAttr, "false">:$issue_token
   );
 
   let assemblyFormat = [{

lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp

@@ -313,7 +313,10 @@ struct DmaToIpuPattern : OpConversionPattern<IpuDmaMemcpyNdOp> {
     // repeat_count
     repeat_count = IntegerAttr::get(i32ty, sizes[3] - 1);
 
-    // issue_token
+    // Set the issue_token
+    issue_token = BoolAttr::get(ctx, op.getIssueToken());
+    // Earlier, all S2MM channels were implicitly assumed to issue a token.
+    // This logic is kept for now for backward compatibility.
     if (!isMM2S)
       issue_token = BoolAttr::get(ctx, true);
 

test/Conversion/DmaToIpu/dma_to_ipu_issue_token.mlir

@@ -0,0 +1,36 @@
+//===- dma_to_ipu_issue_token.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 %s | FileCheck %s
+
+// TODO - more
+// CHECK-LABEL: test1
+// CHECK: aiex.ipu.writebd_shimtile
+// CHECK-SAME: ddr_id = 0 : i32
+// CHECK-SAME: valid_bd = 1 : i32
+// CHECK: aiex.ipu.write32
+// CHECK-SAME: value = 2147483649
+// CHECK: aiex.ipu.writebd_shimtile
+// CHECK-SAME: ddr_id = 1 : i32
+// CHECK: aiex.ipu.write32
+// CHECK-SAME: value = 0
+module  {
+  aie.device(ipu) {
+    memref.global "public" @toMem : memref<16xi32>
+    memref.global "public" @fromMem : memref<16xi32>
+    func.func @test1(%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, issue_token = true } : 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)
+  }
+}