Rebase on main

nod-ai · Oct 9, 2024 · 9d10fd7 · 9d10fd7
1 parent 6f442c2
commit 9d10fd7
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 47 deletions.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaComposition.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaComposition.cpp
@@ -38,54 +38,28 @@ class AMDAIEDmaCompositionPass
 void AMDAIEDmaCompositionPass::runOnOperation() {
   Operation *parentOp = getOperation();
   MLIRContext *context = &getContext();
-
-  auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(parentOp);
-  std::optional<AMDAIEDevice> maybeDevice = getConfigAMDAIEDevice(targetAttr);
-  if (!maybeDevice) {
-    parentOp->emitOpError()
-        << "has no AMDAIEDevice in the target attribute configuration. This "
-           "device-specific information is required to determine when loops "
-           "can be subsumed into DMA operations, and must be attached to a "
-           "containing ModuleOp.";
-    return signalPassFailure();
-  }
-  AMDAIE::AMDAIEDeviceModel deviceModel =
-      AMDAIE::getDeviceModel(maybeDevice.value());
-
-  {
-    RewritePatternSet patterns(context);
-    populateDmaLoopSubsumptionPattern(patterns, std::move(deviceModel),
-                                      onlyZeroStrideOnOuterDim);
-    if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns)))) {
-      parentOp->emitOpError("failed to subsume loops into DMA operations");
-      return signalPassFailure();
-    }
-  }
+  RewritePatternSet patterns(context);
   {
-    RewritePatternSet patterns(context);
-    populateStridedOpCombinationPattern(patterns);
-    if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns)))) {
-      parentOp->emitOpError("failed to compose strided operations");
+    auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(parentOp);
+    std::optional<AMDAIEDevice> maybeDevice = getConfigAMDAIEDevice(targetAttr);
+    if (!maybeDevice) {
+      parentOp->emitOpError()
+          << "has no AMDAIEDevice in the target attribute configuration. This "
+             "device-specific information is required to determine when loops "
+             "can be subsumed into DMA operations, and must be attached to a "
+             "containing ModuleOp.";
       return signalPassFailure();
     }
-  }
-  {
-    RewritePatternSet patterns(context);
-    populateCanonicalizeDoublyStridedOpPatterns(patterns, false);
-    if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns)))) {
-      parentOp->emitOpError(
-          "failed to canonicalize doubly strided DMA operations");
-      return signalPassFailure();
-    }
-  }
-  {
-    RewritePatternSet patterns(context);
+    AMDAIE::AMDAIEDeviceModel deviceModel =
+        AMDAIE::getDeviceModel(maybeDevice.value());
     populateDmaLoopSubsumptionPattern(patterns, std::move(deviceModel),
                                       onlyZeroStrideOnOuterDim);
-    if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns)))) {
-      parentOp->emitOpError("failed to compose strided operations");
-      return signalPassFailure();
-    }
+  }
+  populateStridedOpCombinationPattern(patterns);
+  populateCanonicalizeDoublyStridedOpPatterns(patterns, false);
+  if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns)))) {
+    parentOp->emitOpError("failed to compose strided operations");
+    return signalPassFailure();
   }
 
   IRRewriter rewriter(parentOp->getContext());

diff --git a/...er/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/dma_loop_subsumption_circular.mlir b/...er/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/dma_loop_subsumption_circular.mlir
@@ -77,8 +77,8 @@ module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb}
       amdaie.controlcode {
         scf.forall (%arg2, %arg3) in (2, 2) {
           %1 = affine.apply #map(%arg3)
-          %2 = amdaie.npu.circular_dma_cpy_nd %0([0] [2048] [1], [] [] [])
-          %3 = amdaie.npu.dma_cpy_nd %0([] [] [], [0, %1] [32, 64] [128, 1])
+          amdaie.npu.circular_dma_cpy_nd %0([0] [2048] [1], [] [] [])
+          amdaie.npu.dma_cpy_nd %0([] [] [], [0, %1] [32, 64] [128, 1])
         }
         amdaie.end
       }
@@ -108,8 +108,8 @@ module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb}
       amdaie.controlcode {
         scf.forall (%arg2, %arg3) in (2, 2) {
           %1 = affine.apply #map(%arg3)
-          %2 = amdaie.npu.dma_cpy_nd %0([] [] [], [0, %1] [32, 64] [128, 1])
-          %3 = amdaie.npu.circular_dma_cpy_nd %0([0] [2048] [1], [] [] [])
+          amdaie.npu.dma_cpy_nd %0([] [] [], [0, %1] [32, 64] [128, 1])
+          amdaie.npu.circular_dma_cpy_nd %0([0] [2048] [1], [] [] [])
         }
         amdaie.end
       }