[MLIR] Fix fold-memref-subview-ops for affine.load/store

[bondhugula]

Summary

Fix fold-memref-subview-ops for affine.load/store. We need to expand out
the affine apply on its operands.

[bondhugula]

Rebase.

[tpopp]

Hi,

I am reverting this change as it has a stack-use-after-scope error. I suspect this is due to the call to expandAffineMap returning a SmallVector which is then referenced with the ValueRange indices inside of the if block and then used outside.

[tpopp]

Clang documentation for ASAN, as I feel like everyone has different flows for using it: https://clang.llvm.org/docs/AddressSanitizer.html

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

 //===- FoldSubViewOps.cpp - Fold memref.subview ops -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This transformation pass folds loading/storing from/to subview ops into
 // loading/storing from/to the original memref.
 //
 //===----------------------------------------------------------------------===//
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//

 LogicalResult
 LoadOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy loadOp,
                                              PatternRewriter &rewriter) const {
   auto subViewOp =
       getMemRefOperand(loadOp).template getDefiningOp<memref::SubViewOp>();
   if (!subViewOp)
     return failure();
 
+  ValueRange indices = loadOp.indices();
+  // For affine ops, we need to apply the map to get the operands to get the
+  // "actual" indices.
+  if (auto affineLoadOp = dyn_cast<AffineLoadOp>(loadOp.getOperation())) {
+    auto expandedIndices =
+        expandAffineMap(rewriter, loadOp.getLoc(), affineLoadOp.getAffineMap(),
+                        affineLoadOp.indices());
+    if (!expandedIndices)
+      return failure();
+    indices = expandedIndices.getValue();
+  }
   SmallVector<Value, 4> sourceIndices;
-  if (failed(resolveSourceIndices(loadOp.getLoc(), rewriter, subViewOp,
-                                  loadOp.indices(), sourceIndices)))
+  if (failed(resolveSourceIndices(loadOp.getLoc(), rewriter, subViewOp, indices,
+                                  sourceIndices)))
     return failure();
 
   replaceOp(loadOp, subViewOp, sourceIndices, rewriter);
   return success();
 }
 
 template <typename OpTy>
 LogicalResult
 StoreOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy storeOp,
                                               PatternRewriter &rewriter) const {
   auto subViewOp =
       getMemRefOperand(storeOp).template getDefiningOp<memref::SubViewOp>();
   if (!subViewOp)
     return failure();
 
+  ValueRange indices = storeOp.indices();
+  // For affine ops, we need to apply the map to get the operands to get the
+  // "actual" indices.
+  if (auto affineStoreOp = dyn_cast<AffineStoreOp>(storeOp.getOperation())) {
+    auto expandedIndices =
+        expandAffineMap(rewriter, storeOp.getLoc(),
+                        affineStoreOp.getAffineMap(), affineStoreOp.indices());
+    if (!expandedIndices)
+      return failure();
+    indices = expandedIndices.getValue();
+  }
   SmallVector<Value, 4> sourceIndices;
   if (failed(resolveSourceIndices(storeOp.getLoc(), rewriter, subViewOp,
-                                  storeOp.indices(), sourceIndices)))
+                                  indices, sourceIndices)))
     return failure();
 
   replaceOp(storeOp, subViewOp, sourceIndices, rewriter);
   return success();
 }
 
 void memref::populateFoldSubViewOpPatterns(RewritePatternSet &patterns) {
   patterns.add<LoadOpOfSubViewFolder<AffineLoadOp>,

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

 //  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: %[[C0:.*]] = arith.constant 0 : index
   // 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
+  // Fewer operands than the memref rank.
+  // CHECK-NEXT: affine.apply
+  // CHECK-NEXT: affine.apply #{{.*}}(%[[C0]])[%{{.*}}]
+  // CHECK-NEXT: affine.store
+  affine.store %1, %0[%arg3, 0] : memref<4x4xf32, offset:?, strides: [64, 3]>
   // CHECK-NEXT: return
   return %1 : f32
 }