Move to FoldMemRefAliasOps

Author: lialan

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

@@ -22,9 +22,8 @@ class ConversionTarget;
 namespace amdgpu {
 
 #define GEN_PASS_DECL_AMDGPUEMULATEATOMICSPASS
-#define GEN_PASS_DECL_AMDGPUFOLDSUBVIEWOPSPASS
-#define GEN_PASS_DECL_AMDGPUMASKEDLOADTOLOADPASS
 #define GEN_PASS_DECL_AMDGPURESOLVESTRIDEDMETADATAPASS
+#define GEN_PASS_DECL_AMDGPUMASKEDLOADTOLOADPASS
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/AMDGPU/Transforms/Passes.h.inc"
 
@@ -39,9 +38,6 @@ void populateAmdgpuResolveStridedMetadataPatterns(RewritePatternSet &patterns,
 void populateAmdgpuMaskedloadToLoadPatterns(RewritePatternSet &patterns,
                                             PatternBenefit benefit = 1);
 
-void populateAmdgpuFoldSubviewOpsPatterns(RewritePatternSet &patterns,
-                                          PatternBenefit benefit = 1);
-
 } // namespace amdgpu
 } // namespace mlir
 

mlir/lib/Dialect/AMDGPU/Transforms/CMakeLists.txt

@@ -1,8 +1,7 @@
 add_mlir_dialect_library(MLIRAMDGPUTransforms
   EmulateAtomics.cpp
-  FoldSubviewOps.cpp
-  MaskedloadToLoad.cpp
   ResolveStridedMetadata.cpp
+  MaskedloadToLoad.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/AMDGPU/Transforms

mlir/lib/Dialect/AMDGPU/Transforms/FoldSubviewOps.cpp

@@ -20,6 +20,7 @@
 #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"
@@ -732,6 +733,32 @@ 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 {
+    Location loc = op.getLoc();
+
+    // Check if the source is a subview operation:
+    auto subviewOp = dyn_cast<memref::SubViewOp>(op.getSrc().getDefiningOp());
+    if (!subviewOp)
+      return rewriter.notifyMatchFailure(
+          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);
+
+     rewriter.replaceOpWithNewOp<admgpu::GatherToLDSOp>(
+         op, subviewOp.getSource(), sourceIndices, op.getDst(), op.getDstIndices(),
+         op.getTransferType());
+
+     return success();
+  }
+};
+
 void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
   patterns.add<LoadOpOfSubViewOpFolder<affine::AffineLoadOp>,
                LoadOpOfSubViewOpFolder<memref::LoadOp>,
@@ -762,8 +789,8 @@ void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
                StoreOpOfCollapseShapeOpFolder<memref::StoreOp>,
                StoreOpOfCollapseShapeOpFolder<vector::StoreOp>,
                StoreOpOfCollapseShapeOpFolder<vector::MaskedStoreOp>,
-               SubViewOfSubViewFolder, NVGPUAsyncCopyOpSubViewOpFolder>(
-      patterns.getContext());
+               SubViewOfSubViewFolder, NVGPUAsyncCopyOpSubViewOpFolder,
+               FoldSubviewIntoGatherToLDSOp>(patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//