Support affine.load/store ops in fold-memref-subview-ops pass

[bondhugula]

Summary

Support affine.load/store ops in fold-memref-subview ops pass. The
existing pass just "inlines" the subview operation on load/stores by
inserting affine.apply ops in front of the memref load/store ops: this
is by design always consistent with the semantics on affine.load/store
ops and the same would work even more naturally/intuitively with the
latter. This update leads to a net zero change to lines of code.

[mravishankar]

Neat! Thanks!

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

for (auto index : llvm::seq<size_t>(0, mixedOffsets.size())) {
     Location loc = subViewOp.getLoc();
     sourceIndices[index] = rewriter.create<AffineApplyOp>(
         loc, AffineMap::get(1, numSymbols, expr), dynamicOperands);
   }
   return success();
 }
 
 /// Helpers to access the memref operand for each op.
-static Value getMemRefOperand(memref::LoadOp op) { return op.memref(); }
+template <typename LoadOrStoreOpTy>
+static Value getMemRefOperand(LoadOrStoreOpTy op) {
+  return op.memref();
+}
 
 static Value getMemRefOperand(vector::TransferReadOp op) { return op.source(); }
 
-static Value getMemRefOperand(memref::StoreOp op) { return op.memref(); }
-
 static Value getMemRefOperand(vector::TransferWriteOp op) {
   return op.source();
 }
 
 /// Given the permutation map of the original
 /// `vector.transfer_read`/`vector.transfer_write` operations compute the
 /// permutation map to use after the subview is folded with it.
 static AffineMapAttr getPermutationMapAttr(MLIRContext *context,
LogicalResult matchAndRewrite(OpTy storeOp,
                                 PatternRewriter &rewriter) const override;
 
 private:
   void replaceOp(OpTy storeOp, memref::SubViewOp subViewOp,
                  ArrayRef<Value> sourceIndices,
                  PatternRewriter &rewriter) const;
 };
 
-template <>
-void LoadOpOfSubViewFolder<memref::LoadOp>::replaceOp(
-    memref::LoadOp loadOp, memref::SubViewOp subViewOp,
-    ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
-  rewriter.replaceOpWithNewOp<memref::LoadOp>(loadOp, subViewOp.source(),
-                                              sourceIndices);
+template <typename LoadOpTy>
+void LoadOpOfSubViewFolder<LoadOpTy>::replaceOp(
+    LoadOpTy loadOp, memref::SubViewOp subViewOp, ArrayRef<Value> sourceIndices,
+    PatternRewriter &rewriter) const {
+  rewriter.replaceOpWithNewOp<LoadOpTy>(loadOp, subViewOp.source(),
 }
 
 template <>
 void LoadOpOfSubViewFolder<vector::TransferReadOp>::replaceOp(
     vector::TransferReadOp transferReadOp, memref::SubViewOp subViewOp,
     ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
   // TODO: support 0-d corner case.
   if (transferReadOp.getTransferRank() == 0)
     return;
   rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
       transferReadOp, transferReadOp.getVectorType(), subViewOp.source(),
       sourceIndices,
       getPermutationMapAttr(rewriter.getContext(), subViewOp,
                             transferReadOp.permutation_map()),
       transferReadOp.padding(),
       /*mask=*/Value(), transferReadOp.in_boundsAttr());
 }
 
-template <>
-void StoreOpOfSubViewFolder<memref::StoreOp>::replaceOp(
-    memref::StoreOp storeOp, memref::SubViewOp subViewOp,
+template <typename StoreOpTy>
+void StoreOpOfSubViewFolder<StoreOpTy>::replaceOp(
+    StoreOpTy storeOp, memref::SubViewOp subViewOp,
     ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
-  rewriter.replaceOpWithNewOp<memref::StoreOp>(
-      storeOp, storeOp.value(), subViewOp.source(), sourceIndices);
+  rewriter.replaceOpWithNewOp<StoreOpTy>(storeOp, storeOp.value(),
+                                         subViewOp.source(), sourceIndices);
 }
 
 template <>
 void StoreOpOfSubViewFolder<vector::TransferWriteOp>::replaceOp(
     vector::TransferWriteOp transferWriteOp, memref::SubViewOp subViewOp,
     ArrayRef<Value> sourceIndices, PatternRewriter &rewriter) const {
   // TODO: support 0-d corner case.
   if (transferWriteOp.getTransferRank() == 0)
if (failed(resolveSourceIndices(storeOp.getLoc(), rewriter, subViewOp,
                                   storeOp.indices(), sourceIndices)))
     return failure();
 
   replaceOp(storeOp, subViewOp, sourceIndices, rewriter);
   return success();
 }
 
 void memref::populateFoldSubViewOpPatterns(RewritePatternSet &patterns) {
-  patterns.add<LoadOpOfSubViewFolder<memref::LoadOp>,
+  patterns.add<LoadOpOfSubViewFolder<AffineLoadOp>,
+               LoadOpOfSubViewFolder<memref::LoadOp>,
                LoadOpOfSubViewFolder<vector::TransferReadOp>,
+               StoreOpOfSubViewFolder<AffineStoreOp>,
                StoreOpOfSubViewFolder<memref::StoreOp>,
                StoreOpOfSubViewFolder<vector::TransferWriteOp>>(
       patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/MemRef/fold-subview-ops.mlir

 //  CHECK-SAME:    %[[ARG5:[a-zA-Z0-9]+]]: index
 //  CHECK-SAME:    %[[ARG6:[a-zA-Z0-9]+]]: index
 //  CHECK-SAME:    %[[ARG7:[a-zA-Z0-9]+]]: index
 //   CHECK-DAG:    %[[C0:.+]] = arith.constant 0 : index
 //   CHECK-DAG:    %[[IDX0:.+]] = affine.apply #[[MAP1]](%[[ARG6]])[%[[ARG2]]]
 //   CHECK-DAG:    %[[IDX1:.+]] = affine.apply #[[MAP1]](%[[ARG7]])[%[[ARG3]]]
 //   CHECK-DAG:    vector.transfer_write %[[ARG1]], %[[ARG0]][%[[IDX0]], %[[IDX1]], %[[C0]]]
 //  CHECK-SAME:    {in_bounds = [true], permutation_map = #[[MAP2]]} : vector<4xf32>, memref<?x?x?xf32
+
+// -----
+
+//  Test with affine.load/store ops. We only do a basic test here since the
+//  logic is identical to that with memref.load/store ops. The same affine.apply
+//  ops would be generated.
+
+// CHECK-LABEL: func @fold_static_stride_subview_with_affine_load_store
+func @fold_static_stride_subview_with_affine_load_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> f32 {
+  %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [64, 3]>
+  %1 = affine.load %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [64, 3]>
+  // CHECK-NEXT: affine.apply
+  // CHECK-NEXT: affine.apply
+  // CHECK-NEXT: affine.load
+  affine.store %1, %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [64, 3]>
+  // CHECK-NEXT: affine.apply
+  // CHECK-NEXT: affine.apply
+  // CHECK-NEXT: affine.store
+  // CHECK-NEXT: return
+  return %1 : f32
+}