u[date

nod-ai · Jan 20, 2025 · 7b21be2 · 7b21be2
1 parent 4e21f8d
commit 7b21be2
Show file tree

Hide file tree

Showing 9 changed files with 432 additions and 432 deletions.
diff --git a/...iler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECanonicalizeDoublyStridedOp.cpp b/...iler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECanonicalizeDoublyStridedOp.cpp
@@ -156,12 +156,12 @@ struct FoldDmaOpUnitDims
     SmallVector<OpFoldResult> targetStrides = op.getTargetMixedStrides();
     SmallVector<OpFoldResult> newSourceOffsets, newSourceSizes,
         newSourceStrides, newTargetOffsets, newTargetSizes, newTargetStrides;
-    LogicalResult sourceRes =
-        foldUnitDims(op.getContext(), sourceOffsets, sourceSizes, sourceStrides,
-                     newSourceOffsets, newSourceSizes, newSourceStrides);
-    LogicalResult targetRes =
-        foldUnitDims(op.getContext(), targetOffsets, targetSizes, targetStrides,
-                     newTargetOffsets, newTargetSizes, newTargetStrides);
+    LogicalResult sourceRes = foldUnitDims(
+        *op.getContext(), sourceOffsets, sourceSizes, sourceStrides,
+        newSourceOffsets, newSourceSizes, newSourceStrides);
+    LogicalResult targetRes = foldUnitDims(
+        *op.getContext(), targetOffsets, targetSizes, targetStrides,
+        newTargetOffsets, newTargetSizes, newTargetStrides);
     if (failed(sourceRes) && failed(targetRes)) {
       return failure();
     }

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECombineStridedOps.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECombineStridedOps.cpp
@@ -16,7 +16,6 @@
 #include "iree-amd-aie/Transforms/Transforms.h"
 #include "iree-amd-aie/Transforms/Utils/AMDAIEDmaUtils.h"
 #include "iree-amd-aie/Transforms/Utils/AMDAIEUtils.h"
-#include "llvm/ADT/STLExtras.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #define DEBUG_TYPE "iree-amdaie-combine-strided-ops"
@@ -47,8 +46,10 @@ struct CombineStridedOps
 
     std::unique_ptr<DmaDimConfig> sourceDmaDimConfig;
     std::unique_ptr<DmaDimConfig> targetDmaDimConfig;
+
     SmallVector<Operation *> userOpsToBeErased;
     AMDAIE::DoublyStridedOpInterface nextStridedOp;
+
     if (auto npuDmaOp = dyn_cast<AMDAIE::NpuDmaCpyNdOp>(op.getOperation())) {
       LLVM_DEBUG(llvm::dbgs() << "npuDmaOp: " << npuDmaOp << "\n");
       // Fail if any non-wait user operations.
@@ -105,6 +106,10 @@ struct CombineStridedOps
       return failure();
     }
 
+    MLIRContext &context = *rewriter.getContext();
+    auto dimCountCheck = std::bind(&DmaDimConfig::exceedsNbDims,
+                                   std::ref(sourceDmaDimConfig), _1);
+
     SmallVector<OpFoldResult> sourceOffsetsA = op.getSourceMixedOffsets();
     SmallVector<OpFoldResult> sourceSizesA = op.getSourceMixedSizes();
     SmallVector<OpFoldResult> sourceStridesA = op.getSourceMixedStrides();
@@ -114,11 +119,15 @@ struct CombineStridedOps
         nextStridedOp.getSourceMixedSizes();
     SmallVector<OpFoldResult> sourceStridesB =
         nextStridedOp.getSourceMixedStrides();
-    bool areSourcesCombinable = areAccessPatternsCombinable(
-        sourceOffsetsA, sourceSizesA, sourceStridesA, sourceOffsetsB,
-        sourceSizesB, sourceStridesB,
-        std::bind(&DmaDimConfig::exceedsNbDims, std::ref(sourceDmaDimConfig),
-                  _1));
+    SmallVector<OpFoldResult> newSourceOffsets;
+    SmallVector<OpFoldResult> newSourceSizes;
+    SmallVector<OpFoldResult> newSourceStrides;
+    if (failed(combineAccessPatterns(
+            context, sourceOffsetsA, sourceSizesA, sourceStridesA,
+            sourceOffsetsB, sourceSizesB, sourceStridesB, newSourceOffsets,
+            newSourceSizes, newSourceStrides, dimCountCheck))) {
+      return failure();
+    }
 
     SmallVector<OpFoldResult> targetOffsetsA = op.getTargetMixedOffsets();
     SmallVector<OpFoldResult> targetSizesA = op.getTargetMixedSizes();
@@ -129,53 +138,25 @@ struct CombineStridedOps
         nextStridedOp.getTargetMixedSizes();
     SmallVector<OpFoldResult> targetStridesB =
         nextStridedOp.getTargetMixedStrides();
-    bool areTargetsCombinable = areAccessPatternsCombinable(
-        targetOffsetsA, targetSizesA, targetStridesA, targetOffsetsB,
-        targetSizesB, targetStridesB,
-        std::bind(&DmaDimConfig::exceedsNbDims, std::ref(targetDmaDimConfig),
-                  _1));
-
-    LLVM_DEBUG(llvm::dbgs()
-               << "areSourcesCombinable: " << areSourcesCombinable << "\n");
-    LLVM_DEBUG(llvm::dbgs()
-               << "areTargetsCombinable: " << areTargetsCombinable << "\n");
-
-    if (areSourcesCombinable && areTargetsCombinable) {
-      SmallVector<OpFoldResult> newSourceOffsets;
-      SmallVector<OpFoldResult> newSourceSizes;
-      SmallVector<OpFoldResult> newSourceStrides;
-      if (failed(combineAccessPatterns(
-              rewriter, sourceOffsetsA, sourceSizesA, sourceStridesA,
-              sourceOffsetsB, sourceSizesB, sourceStridesB, newSourceOffsets,
-              newSourceSizes, newSourceStrides,
-              std::bind(&DmaDimConfig::exceedsNbDims,
-                        std::ref(sourceDmaDimConfig), _1)))) {
-        return failure();
-      }
-
-      SmallVector<OpFoldResult> newTargetOffsets;
-      SmallVector<OpFoldResult> newTargetSizes;
-      SmallVector<OpFoldResult> newTargetStrides;
-      if (failed(combineAccessPatterns(
-              rewriter, targetOffsetsA, targetSizesA, targetStridesA,
-              targetOffsetsB, targetSizesB, targetStridesB, newTargetOffsets,
-              newTargetSizes, newTargetStrides,
-              std::bind(&DmaDimConfig::exceedsNbDims,
-                        std::ref(targetDmaDimConfig), _1)))) {
-        return failure();
-      }
+    SmallVector<OpFoldResult> newTargetOffsets;
+    SmallVector<OpFoldResult> newTargetSizes;
+    SmallVector<OpFoldResult> newTargetStrides;
+    if (failed(combineAccessPatterns(
+            context, targetOffsetsA, targetSizesA, targetStridesA,
+            targetOffsetsB, targetSizesB, targetStridesB, newTargetOffsets,
+            newTargetSizes, newTargetStrides, dimCountCheck))) {
+      return failure();
+    }
 
-      rewriter.setInsertionPoint(op);
-      auto newDoublyStridedOp = nextStridedOp.createDoublyStridedOp(
-          rewriter, newTargetOffsets, newTargetSizes, newTargetStrides,
-          newSourceOffsets, newSourceSizes, newSourceStrides);
-      rewriter.replaceOp(nextStridedOp, newDoublyStridedOp.getOperation());
+    rewriter.setInsertionPoint(op);
+    auto newDoublyStridedOp = nextStridedOp.createDoublyStridedOp(
+        rewriter, newTargetOffsets, newTargetSizes, newTargetStrides,
+        newSourceOffsets, newSourceSizes, newSourceStrides);
+    rewriter.replaceOp(nextStridedOp, newDoublyStridedOp.getOperation());
 
-      for (Operation *userOp : userOpsToBeErased) rewriter.eraseOp(userOp);
-      rewriter.eraseOp(op);
-      return success();
-    }
-    return failure();
+    for (Operation *userOp : userOpsToBeErased) rewriter.eraseOp(userOp);
+    rewriter.eraseOp(op);
+    return success();
   }
 
   template <typename T>

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
@@ -55,7 +55,6 @@ void AMDAIEDmaCompositionPass::runOnOperation() {
                                     onlyZeroStrideOnOuterDim);
   populateStridedOpCombinationPattern(patterns);
   populateCanonicalizeDoublyStridedOpPatterns(patterns, false, deviceModel);
-
   if (failed(applyPatternsGreedily(parentOp, std::move(patterns)))) {
     parentOp->emitOpError("failed to compose strided operations");
     return signalPassFailure();

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELowerToAIE.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELowerToAIE.cpp
@@ -345,7 +345,7 @@ LogicalResult AIEDeviceBuilder::foldDimsAndReturnAsStatic(
   SmallVector<OpFoldResult> offsets(
       strides.size(), getAsIndexOpFoldResult(rewriter.getContext(), 0));
   SmallVector<OpFoldResult> unitOffsets, unitSizes, unitStrides, newOffsets;
-  (void)foldUnitDims(rewriter.getContext(), offsets, sizes, strides,
+  (void)foldUnitDims(*rewriter.getContext(), offsets, sizes, strides,
                      unitOffsets, unitSizes, unitStrides);
   DmaDimConfig dmaDimConfig(deviceModel, memSpace);
   SmallVector<int64_t> maxSizes = dmaDimConfig.getMaxSizes(unitOffsets.size());