Still ... merge it with FoldMemRefAliasOps pass.

Author: lialan

mlir/include/mlir/Dialect/AMDGPU/Transforms/Passes.td

@@ -70,16 +70,4 @@ def AmdgpuMaskedloadToLoadPass : Pass<"amdgpu-maskedload-to-load"> {
     "memref::MemRefDialect"
   ];
 }
-
-def AmdgpuFoldSubviewOpsPass : Pass<"amdgpu-fold-subview-ops"> {
-  let summary = "Fold subview operations into their parent operations";
-  let description = [{
-    This pass identifies `memref.subview` sources of `GatherToLDSOp` and
-    attempts to fold the source ops, potentially simplifying the overall
-    operation and improving performance.
-  }];
-  let dependentDialects = [
-    "memref::MemRefDialect"
-  ];
-}
 #endif // MLIR_DIALECT_AMDGPU_TRANSFORMS_PASSES_TD_

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

@@ -11,6 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
@@ -20,7 +21,6 @@
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
-#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -733,11 +733,11 @@ LogicalResult NVGPUAsyncCopyOpSubViewOpFolder::matchAndRewrite(
   return success();
 }
 
-struct FoldSubviewIntoGatherToLDSOp
-    : public OpRewritePattern<amdgpu::GatherToLDSOp> {
-  using OpRewritePattern<amdgpu::GatherToLDSOp>::OpRewritePattern;
-  LogicalResult
-  matchAndRewrite(amdgpu::GatherToLDSOp op, PatternRewriter &rewriter) const override {
+struct FoldSubviewIntoAMDGPUGatherToLDSOp final
+    : OpRewritePattern<amdgpu::GatherToLDSOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(amdgpu::GatherToLDSOp op,
+                                PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
 
     // Check if the source is a subview operation:
@@ -747,15 +747,15 @@ struct FoldSubviewIntoGatherToLDSOp
           loc, "GatherToLDSOp can only be folded if the source is a SubviewOp");
 
     SmallVector<Value> sourceIndices;
-     mlir::affine::resolveIndicesIntoOpWithOffsetsAndStrides(
-         rewriter, loc, subviewOp.getMixedOffsets(), subviewOp.getMixedStrides(),
-         subviewOp.getDroppedDims(), op.getSrcIndices(), sourceIndices);
+    mlir::affine::resolveIndicesIntoOpWithOffsetsAndStrides(
+        rewriter, loc, subviewOp.getMixedOffsets(), subviewOp.getMixedStrides(),
+        subviewOp.getDroppedDims(), op.getSrcIndices(), sourceIndices);
 
-     rewriter.replaceOpWithNewOp<admgpu::GatherToLDSOp>(
-         op, subviewOp.getSource(), sourceIndices, op.getDst(), op.getDstIndices(),
-         op.getTransferType());
+    rewriter.replaceOpWithNewOp<amdgpu::GatherToLDSOp>(
+        op, subviewOp.getSource(), sourceIndices, op.getDst(),
+        op.getDstIndices(), op.getTransferType());
 
-     return success();
+    return success();
   }
 };
 
@@ -790,7 +790,7 @@ void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
                StoreOpOfCollapseShapeOpFolder<vector::StoreOp>,
                StoreOpOfCollapseShapeOpFolder<vector::MaskedStoreOp>,
                SubViewOfSubViewFolder, NVGPUAsyncCopyOpSubViewOpFolder,
-               FoldSubviewIntoGatherToLDSOp>(patterns.getContext());
+               FoldSubviewIntoAMDGPUGatherToLDSOp>(patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/AMDGPU/amdgpu-fold-subviews.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -amdgpu-fold-subview-ops -split-input-file %s | FileCheck %s
+// RUN: mlir-opt --fold-memref-alias-ops --split-input-file %s | FileCheck %s
 
 #gpu_lds_addrspace = 3