[mlir][memref] Fold subview(subview(x))

[springerm]

Summary

Folding of rank-reduced subviews is also supported.

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

 class StoreOpOfCollapseShapeOpFolder final : public OpRewritePattern<OpTy> {
 public:
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy storeOp,
                                 PatternRewriter &rewriter) const override;
 };
 
+/// Folds subview(subview(x)) to a single subview(x).
+class SubViewOfSubViewFolder : public OpRewritePattern<memref::SubViewOp> {
+public:
+  using OpRewritePattern<memref::SubViewOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::SubViewOp subView,
+                                PatternRewriter &rewriter) const override {
+    Location loc = subView.getLoc();
+    auto srcSubView = subView.getSource().getDefiningOp<memref::SubViewOp>();
+    if (!srcSubView)
+      return failure();
+    int64_t srcRank = srcSubView.getSourceType().getRank();
+
+    // TODO: Only stride 1 is supported.
+    for (auto s : {subView.getMixedStrides(), srcSubView.getMixedStrides()})
+      if (!llvm::all_of(
+              s, [](OpFoldResult ofr) { return isConstantIntValue(ofr, 1); }))
+        return failure();
+
+    // Get original offsets and sizes.
+    SmallVector<OpFoldResult> offsets = subView.getMixedOffsets();
+    SmallVector<OpFoldResult> srcOffsets = srcSubView.getMixedOffsets();
+    SmallVector<OpFoldResult> sizes = subView.getMixedSizes();
+    SmallVector<OpFoldResult> srcSizes = srcSubView.getMixedSizes();
+
+    // Compute new offsets and sizes.
+    llvm::SmallBitVector srcReducedDims = srcSubView.getDroppedDims();
+    SmallVector<OpFoldResult> newOffsets, newSizes;
+    int64_t dim = 0;
+    for (int64_t srcDim = 0; srcDim < srcRank; ++srcDim) {
+      if (srcReducedDims[srcDim]) {
+        // Dim is reduced in srcSubView.
+        assert(isConstantIntValue(srcSizes[srcDim], 1) && "expected size 1");
+        newOffsets.push_back(srcOffsets[srcDim]);
+        newSizes.push_back(srcSizes[srcDim]);
+        continue;
+      }
+      AffineExpr sym0, sym1;
+      bindSymbols(subView.getContext(), sym0, sym1);
+      newOffsets.push_back(makeComposedFoldedAffineApply(
+          rewriter, loc, sym0 + sym1, {srcOffsets[srcDim], offsets[dim]}));
+      newSizes.push_back(sizes[dim]);
+      ++dim;
+    }
+
+    // Replace original op.
+    rewriter.replaceOpWithNewOp<memref::SubViewOp>(
+        subView, subView.getType(), srcSubView.getSource(), newOffsets,
+        newSizes, srcSubView.getMixedStrides());
+    return success();
+  }
+};
 } // namespace
 
 static SmallVector<Value>
 calculateExpandedAccessIndices(AffineMap affineMap,
                                const SmallVector<Value> &indices, Location loc,
                                PatternRewriter &rewriter) {
   SmallVector<Value> expandedIndices;
   for (unsigned i = 0, e = affineMap.getNumResults(); i < e; i++)
patterns.add<LoadOpOfSubViewOpFolder<AffineLoadOp>,
                StoreOpOfSubViewOpFolder<vector::TransferWriteOp>,
                LoadOpOfExpandShapeOpFolder<AffineLoadOp>,
                LoadOpOfExpandShapeOpFolder<memref::LoadOp>,
                StoreOpOfExpandShapeOpFolder<AffineStoreOp>,
                StoreOpOfExpandShapeOpFolder<memref::StoreOp>,
                LoadOpOfCollapseShapeOpFolder<AffineLoadOp>,
                LoadOpOfCollapseShapeOpFolder<memref::LoadOp>,
                StoreOpOfCollapseShapeOpFolder<AffineStoreOp>,
-               StoreOpOfCollapseShapeOpFolder<memref::StoreOp>>(
-      patterns.getContext());
+               StoreOpOfCollapseShapeOpFolder<memref::StoreOp>,
+               SubViewOfSubViewFolder>(patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir

-// RUN: mlir-opt -fold-memref-alias-ops -split-input-file %s -o - | FileCheck %s
+// RUN: mlir-opt -fold-memref-alias-ops -split-input-file %s | FileCheck %s
 
 func.func @fold_static_stride_subview_with_load(%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, strided<[64, 3], offset: ?>>
   %1 = memref.load %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>
   return %1 : f32
 }
 //  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
 //  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (d0 * 3 + s0)>
affine.for %arg2 = 0 to 3 {
     %1 = affine.load %0[] : memref<f32>
     affine.store %1, %arg1[0] : memref<1xf32>
   }
   return %arg1 : memref<1xf32>
 }
 // CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index
 // CHECK-NEXT: affine.for %{{.*}} = 0 to 3 {
 // CHECK-NEXT:   affine.load %[[ARG0]][%[[ZERO]]] : memref<1xf32>
+
+// -----
+
+//       CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 2)>
+// CHECK-LABEL: func @subview_of_subview(
+//  CHECK-SAME:     %[[m:.*]]: memref<1x1024xf32, 3>, %[[pos:.*]]: index
+//       CHECK:   %[[add:.*]] = affine.apply #[[$map]]()[%arg1]
+//       CHECK:   memref.subview %arg0[4, %[[add]]] [1, 1] [1, 1] : memref<1x1024xf32, 3> to memref<f32, strided<[], offset: ?>, 3>
+func.func @subview_of_subview(%m: memref<1x1024xf32, 3>, %pos: index)
+    -> memref<f32, strided<[], offset: ?>, 3>
+{
+  %0 = memref.subview %m[3, %pos] [1, 2] [1, 1]
+      : memref<1x1024xf32, 3>
+        to memref<1x2xf32, strided<[1024, 2], offset: ?>, 3>
+  %1 = memref.subview %0[1, 2] [1, 1] [1, 1]
+      : memref<1x2xf32, strided<[1024, 2], offset: ?>, 3>
+        to memref<f32, strided<[], offset: ?>, 3>
+  return %1 : memref<f32, strided<[], offset: ?>, 3>
+}
+
+// -----
+
+// CHECK-LABEL: func @subview_of_subview_rank_reducing(
+//  CHECK-SAME:     %[[m:.*]]: memref<?x?x?xf32>
+//       CHECK:   memref.subview %arg0[3, 7, 8] [1, 1, 1] [1, 1, 1] : memref<?x?x?xf32> to memref<f32, strided<[], offset: ?>>
+func.func @subview_of_subview_rank_reducing(%m: memref<?x?x?xf32>,
+                                            %sz: index, %pos: index)
+    -> memref<f32, strided<[], offset: ?>>
+{
+  %0 = memref.subview %m[3, 1, 8] [1, %sz, 1] [1, 1, 1]
+      : memref<?x?x?xf32>
+        to memref<?xf32, strided<[1], offset: ?>>
+  %1 = memref.subview %0[6] [1] [1]
+      : memref<?xf32, strided<[1], offset: ?>>
+        to memref<f32, strided<[], offset: ?>>
+  return %1 : memref<f32, strided<[], offset: ?>>
+}