[DmaComposition] Rework logic to enable new composition opportunity

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECanonicalizeDoublyStridedOp.cpp

@@ -154,22 +154,16 @@ struct FoldDmaOpUnitDims
     SmallVector<OpFoldResult> targetOffsets = op.getTargetMixedOffsets();
     SmallVector<OpFoldResult> targetSizes = op.getTargetMixedSizes();
     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);
-    if (failed(sourceRes) && failed(targetRes)) {
-      return failure();
-    }
+    LogicalResult sourceRes = foldUnitDims(op.getContext(), sourceOffsets,
+                                           sourceSizes, sourceStrides);
+    LogicalResult targetRes = foldUnitDims(op.getContext(), targetOffsets,
+                                           targetSizes, targetStrides);
+    if (failed(sourceRes) && failed(targetRes)) return failure();
 
     rewriter.setInsertionPointAfter(op);
     auto newDoublyStridedOp = op.createDoublyStridedOp(
-        rewriter, newTargetOffsets, newTargetSizes, newTargetStrides,
-        newSourceOffsets, newSourceSizes, newSourceStrides);
+        rewriter, targetOffsets, targetSizes, targetStrides, sourceOffsets,
+        sourceSizes, sourceStrides);
     rewriter.replaceOp(op, newDoublyStridedOp.getOperation());
     return success();
   }

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 *ctx = 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(
+            ctx, 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(
+            ctx, 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>

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELowerToAIE.cpp

@@ -28,7 +28,6 @@
 #include "llvm/Support/Debug.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/IRMapping.h"
-#include "mlir/IR/Iterators.h"
 #include "mlir/Pass/PassManager.h"
 
 #define DEBUG_TYPE "iree-amdaie-lower-to-aie"
@@ -344,14 +343,13 @@ 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,
-                     unitOffsets, unitSizes, unitStrides);
+  (void)foldUnitDims(rewriter.getContext(), offsets, sizes, strides);
+
   DmaDimConfig dmaDimConfig(deviceModel, memSpace);
-  SmallVector<int64_t> maxSizes = dmaDimConfig.getMaxSizes(unitOffsets.size());
+  SmallVector<int64_t> maxSizes = dmaDimConfig.getMaxSizes(offsets.size());
   SmallVector<OpFoldResult> linearOffsets, linearSizes, linearStrides;
   (void)foldLinearDims(
-      rewriter.getContext(), unitOffsets, unitSizes, unitStrides, linearOffsets,
+      rewriter.getContext(), offsets, sizes, strides, linearOffsets,
       linearSizes, linearStrides, [&](size_t idxFromEnd, int64_t size) {
         return idxFromEnd < maxSizes.size() &&
                size <= maxSizes[maxSizes.size() - idxFromEnd - 1];