From 066ce8f2224deb81141ff1140771ae95864ae7fe Mon Sep 17 00:00:00 2001
From: Jorn Tuyls <jtuyls@users.noreply.github.com>
Date: Wed, 10 Apr 2024 23:08:20 +0200
Subject: [PATCH] Add `issue_token` attribute to IpuDmaMemcpyNdOp (#1209)

Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
---
 include/aie/Dialect/AIEX/IR/AIEX.td           | 13 +++++--
 lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp   |  5 ++-
 .../DmaToIpu/dma_to_ipu_issue_token.mlir      | 36 +++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 test/Conversion/DmaToIpu/dma_to_ipu_issue_token.mlir

diff --git a/include/aie/Dialect/AIEX/IR/AIEX.td b/include/aie/Dialect/AIEX/IR/AIEX.td
index dc399d9e03..11057e8765 100644
--- a/include/aie/Dialect/AIEX/IR/AIEX.td
+++ b/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 = [{
diff --git a/lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp b/lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp
index 8c57a35343..4d68d086d0 100644
--- a/lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp
+++ b/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);
 
diff --git a/test/Conversion/DmaToIpu/dma_to_ipu_issue_token.mlir b/test/Conversion/DmaToIpu/dma_to_ipu_issue_token.mlir
new file mode 100644
index 0000000000..4eb5b02bdc
--- /dev/null
+++ b/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)
+  }
+}