for test

nod-ai · Sep 23, 2024 · e924000 · e924000
1 parent 326b930
commit e924000
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/...ler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIETransferStridedAccessPattern.cpp b/...ler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIETransferStridedAccessPattern.cpp
@@ -240,6 +240,7 @@ static LogicalResult createNewAddressing(
 static LogicalResult transferDmaAddressing(MLIRContext *ctx,
                                            AMDAIE::ConnectionOp connectionOp) {
   IRRewriter rewriter(ctx);
+  OpBuilder::InsertionGuard guard(rewriter);
 
   FailureOr<AMDAIE::NpuCircularDmaCpyNdOp> maybeNpuDmaUserOp =
       connectionOp.getNpuCircularDmaCpyNdUser();
@@ -265,6 +266,7 @@ static LogicalResult transferDmaAddressing(MLIRContext *ctx,
   // Change the source/target addressing of all users from a connection op.
   for (Operation *user : connectionOp->getUsers()) {
     if (auto dmaOp = dyn_cast<AMDAIE::NpuDmaCpyNdOp>(user)) {
+      OpBuilder::InsertionGuard guard(rewriter);
       SmallVector<OpFoldResult> srcOffsets = dmaOp.getSourceMixedOffsets();
       SmallVector<OpFoldResult> srcSizes = dmaOp.getSourceMixedSizes();
       SmallVector<OpFoldResult> srcStrides = dmaOp.getSourceMixedStrides();

diff --git a/.../plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/transfer_strided_access_pattern.mlir b/.../plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/transfer_strided_access_pattern.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-amdaie-transfer-strided-access-pattern))" --split-input-file %s --verify-diagnostics | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-amdaie-transfer-strided-access-pattern))" --split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: @single_dma_l3_source
 // CHECK:       %[[APPLY:.+]] = affine.apply
@@ -58,6 +58,9 @@ module {
 #map = affine_map<(d0) -> (d0 * 64)>
 module {
   func.func @multiple_dma_l3_source(%arg0: !amdaie.logicalobjectfifo<memref<2048xi32, 1 : i32>, 2>, %arg1: !amdaie.logicalobjectfifo<memref<256x128xi32>>) {
+    %c1 = arith.constant 1 : index
+    %c0 = arith.constant 0 : index
+    %c2 = arith.constant 2 : index
     amdaie.workgroup {
       %0 = amdaie.connection(%arg0, %arg1) : (!amdaie.logicalobjectfifo<memref<2048xi32, 1 : i32>, 2>, !amdaie.logicalobjectfifo<memref<256x128xi32>>)
       amdaie.controlcode {
@@ -66,8 +69,10 @@ module {
           %2 = amdaie.npu.circular_dma_cpy_nd %0([0] [2048] [1], [] [] [])
           %3 = amdaie.npu.dma_cpy_nd %0([] [] [], [0, 0, %1] [2, 32, 32] [32, 128, 1])
           amdaie.npu.dma_wait(%3, MM2S)
-          %4 = amdaie.npu.dma_cpy_nd %0([] [] [], [0, 224, %1] [2, 32, 32] [32, 128, 1])
-          amdaie.npu.dma_wait(%4, MM2S)
+          scf.for %arg4 = %c0 to %c2 step %c1 {
+            %4 = amdaie.npu.dma_cpy_nd %0([] [] [], [0, 224, %1] [2, 32, 32] [32, 128, 1])
+            amdaie.npu.dma_wait(%4, MM2S)
+          }
         }
         amdaie.end
       }