[mlir][memref] Propagate LoadOp/StoreOp to SubViewOp source where possible

[Hardcode84]

Summary

• Rank-reducing subview is not supported yet
• Subview offsets must be statically known to be 0s and strides must 1s

[nicolasvasilache]

Thanks @Hardcode84 !
Note that these patterns cannot remain canonicalizations in general because they will have far reaching consequences.
Can you put them in a new pass?

[Hardcode84]

And what exactly wrong with making them canonicalizations?

[nicolasvasilache]

They fold away subviews indiscriminately as soon as canonicalization is invoked.
This prematurely drops information that other parts of the compiler want to operate on.
Actually I think @mravishankar implemented similar functionality to fold subview away in the past, not sure whether they made it into core or just IREE?

[Hardcode84]

Actually, there is a FoldSubViewOpsPass, but it can generate additional ops, which can be undesirable during canonicalization. Here I specifically restricted it in such way it never generates additional ops.

[Hardcode84]

So, what useful information we are losing here? Subview dimensions? Is there any examples they can be usefull in this context?

The only example I can see is someone doing bounds-checked array accesses and doing this very late in pipeline (at memref->llvm lowering level), but this not very good idea in general IMO, regardless of this specific transformation.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

Value newSubView = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), resultType, castOp.source(), subViewOp.offsets(),
         subViewOp.sizes(), subViewOp.strides(), subViewOp.static_offsets(),
         subViewOp.static_sizes(), subViewOp.static_strides());
     rewriter.replaceOpWithNewOp<CastOp>(subViewOp, subViewOp.getType(),
                                         newSubView);
     return success();
   }
 };
-} // namespace
 
 /// Return the canonical type of the result of a subview.
 struct SubViewReturnTypeCanonicalizer {
   MemRefType operator()(SubViewOp op, ArrayRef<OpFoldResult> mixedOffsets,
                         ArrayRef<OpFoldResult> mixedSizes,
                         ArrayRef<OpFoldResult> mixedStrides) {
     return getCanonicalSubViewResultType(op.getType().getRank(),
                                          op.getSourceType(), mixedOffsets,
                                          mixedSizes, mixedStrides);
   }
 };
 
 /// A canonicalizer wrapper to replace SubViewOps.
 struct SubViewCanonicalizer {
   void operator()(PatternRewriter &rewriter, SubViewOp op, SubViewOp newOp) {
     rewriter.replaceOpWithNewOp<CastOp>(op, newOp, op.getType());
   }
 };
 
+static bool isSameRank(Type type1, Type type2) {
+  auto shaped1 = type1.dyn_cast<ShapedType>();
+  if (!shaped1)
+    return false;
+
+  auto shaped2 = type2.dyn_cast<ShapedType>();
+  if (!shaped2)
+    return false;
+
+  if (!shaped1.hasRank() || !shaped2.hasRank())
+    return false;
+
+  return shaped1.getRank() == shaped2.getRank();
+}
+
+static bool isMixedValuesEqual(ArrayRef<OpFoldResult> values,
+                               int64_t expectedVal) {
+  for (auto val : values) {
+    auto intVal = getConstantIntValue(val);
+    if (!intVal || *intVal != expectedVal)
+      return false;
+  }
+  return true;
+}
+
+struct SubviewLoadPropagate : public OpRewritePattern<LoadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(LoadOp op,
+                                PatternRewriter &rewriter) const override {
+    auto src = op.memref().getDefiningOp<SubViewOp>();
+    if (!src)
+      return failure();
+
+    if (!isSameRank(src.source().getType(), src.getType()))
+      return failure();
+
+    if (!isMixedValuesEqual(src.getMixedOffsets(), 0) ||
+        !isMixedValuesEqual(src.getMixedStrides(), 1))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<LoadOp>(op, src.source(), op.indices());
+    return success();
+  }
+};
+
+struct SubviewStorePropagate : public OpRewritePattern<StoreOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(StoreOp op,
+                                PatternRewriter &rewriter) const override {
+    auto src = op.memref().getDefiningOp<SubViewOp>();
+    if (!src)
+      return failure();
+
+    if (!isSameRank(src.source().getType(), src.getType()))
+      return failure();
+
+    if (!isMixedValuesEqual(src.getMixedOffsets(), 0) ||
+        !isMixedValuesEqual(src.getMixedStrides(), 1))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<StoreOp>(op, op.value(), src.source(),
+                                         op.indices());
+    return success();
+  }
+};
+} // namespace
+
 void SubViewOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results
       .add<OpWithOffsetSizesAndStridesConstantArgumentFolder<
                SubViewOp, SubViewReturnTypeCanonicalizer, SubViewCanonicalizer>,
-           SubViewOpMemRefCastFolder>(context);
+           SubViewOpMemRefCastFolder, SubviewLoadPropagate,
+           SubviewStorePropagate>(context);
 }
 
 OpFoldResult SubViewOp::fold(ArrayRef<Attribute> operands) {
   auto resultShapedType = getResult().getType().cast<ShapedType>();
   auto sourceShapedType = source().getType().cast<ShapedType>();
 
   if (resultShapedType.hasStaticShape() &&
       resultShapedType == sourceShapedType) {

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK-LABEL: func @allocator
 // CHECK:   %[[alloc:.+]] = memref.alloc
 // CHECK:   memref.store %[[alloc:.+]], %arg0
 func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index)  {
   %0 = memref.alloc(%arg1) : memref<?xi32>
   memref.store %0, %arg0[] : memref<memref<?xi32>>
   return 
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+func @subview_load_canonicalize(%arg0 : memref<?x?x?xf32>, %arg1 : index, %arg2 : index) -> f32
+{
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c4 = constant 4 : index
+  %0 = memref.subview %arg0[%c0, %c0, %c0] [%c4, %c1, %arg1] [%c1, %c1, %c1] : memref<?x?x?xf32> to memref<?x?x?xf32, #map0>
+  %1 = memref.load %0[%arg2, %c0, %c1] : memref<?x?x?xf32, #map0>
+  return %1 : f32
+}
+// CHECK-LABEL: func @subview_load_canonicalize
+//  CHECK-SAME:   %[[ARG0:.+]]: memref<?x?x?xf32>, %[[ARG1:.+]]: index, %[[ARG2:.+]]: index
+//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//       CHECK:   %[[RESULT:.+]] = memref.load %[[ARG0]][%[[ARG2]], %[[C0]], %[[C1]]]
+//       CHEKC:   return %[[RESULT]]
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+func @subview_store_canonicalize(%arg0 : memref<?x?x?xf32>, %arg1 : index, %arg2 : index)
+{
+  %cst42 = constant 0.0 : f32
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c4 = constant 4 : index
+  %0 = memref.subview %arg0[%c0, %c0, %c0] [%c4, %c1, %arg1] [%c1, %c1, %c1] : memref<?x?x?xf32> to memref<?x?x?xf32, #map0>
+  memref.store %cst42, %0[%arg2, %c0, %c1] : memref<?x?x?xf32, #map0>
+  return
+}
+// CHECK-LABEL: func @subview_store_canonicalize
+//  CHECK-SAME:   %[[ARG0:.+]]: memref<?x?x?xf32>, %[[ARG1:.+]]: index, %[[ARG2:.+]]: index
+//   CHECK-DAG:   %[[VAL:.+]] = constant 0.000000e+00 : f32
+//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//       CHECK:   memref.store %[[VAL]], %[[ARG0]][%[[ARG2]], %[[C0]], %[[C1]]]
+//       CHEKC:   return

mlir/test/Transforms/canonicalize.mlir

func @view(%arg0 : index) -> (f32, f32, f32, f32) {
   %4 = memref.cast %0 : memref<2048xi8> to memref<?xi8>
   %5 = memref.view %4[%c15][%c15, %c7] : memref<?xi8> to memref<?x?xf32>
   %r3 = memref.load %5[%c0, %c0] : memref<?x?xf32>
   return %r0, %r1, %r2, %r3 : f32, f32, f32, f32
 }
 
 // -----
 
+// Function to avoid load/store folding to subview source
+func private @break_chain(%arg0: memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>) -> memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
+
 // CHECK-DAG: #[[$BASE_MAP0:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>
 // CHECK-DAG: #[[$SUBVIEW_MAP0:map[0-9]+]] = affine_map<(d0, d1, d2)[s0] -> (d0 * 64 + s0 + d1 * 4 + d2)>
 // CHECK-DAG: #[[$SUBVIEW_MAP1:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2 + 79)>
 // CHECK-DAG: #[[$SUBVIEW_MAP2:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 128 + d1 * 28 + d2 * 11)>
 // CHECK-DAG: #[[$SUBVIEW_MAP3:map[0-9]+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
 // CHECK-DAG: #[[$SUBVIEW_MAP4:map[0-9]+]] = affine_map<(d0, d1, d2)[s0] -> (d0 * 128 + s0 + d1 * 28 + d2 * 11)>
 // CHECK-DAG: #[[$SUBVIEW_MAP5:map[0-9]+]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + d2 * s2 + 79)>
 // CHECK-DAG: #[[$SUBVIEW_MAP6:map[0-9]+]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2 * 2)>
func @subview(%arg0 : index, %arg1 : index) -> (index, index) {
   // Note that the subview uses the base memrefs layout map because it used
   // zero offset and unit stride arguments.
   // CHECK: memref.subview %[[ALLOC0]][0, 0, 0] [7, 11, 2] [1, 1, 1] :
   // CHECK-SAME: memref<8x16x4xf32, #[[$BASE_MAP0]]>
   // CHECK-SAME: to memref<7x11x2xf32, #[[$BASE_MAP0]]>
   %1 = memref.subview %0[%c0, %c0, %c0] [%c7, %c11, %c2] [%c1, %c1, %c1]
     : memref<8x16x4xf32, offset : 0, strides : [64, 4, 1]> to
       memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
-  %v0 = memref.load %1[%c0, %c0, %c0] : memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
+  %temp1 = call @break_chain(%1) : (memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>) -> (memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>)
+  %v0 = memref.load %temp1[%c0, %c0, %c0] : memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
 
   // Test: subview with one dynamic operand can also be folded.
   // CHECK: memref.subview %[[ALLOC0]][0, %[[ARG0]], 0] [7, 11, 15] [1, 1, 1] :
   // CHECK-SAME: memref<8x16x4xf32, #[[$BASE_MAP0]]>
   // CHECK-SAME: to memref<7x11x15xf32, #[[$SUBVIEW_MAP0]]>
   %2 = memref.subview %0[%c0, %arg0, %c0] [%c7, %c11, %c15] [%c1, %c1, %c1]
     : memref<8x16x4xf32, offset : 0, strides : [64, 4, 1]> to
       memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
-  memref.store %v0, %2[%c0, %c0, %c0] : memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
+  %temp2 = call @break_chain(%2) : (memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>) -> (memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>)
+  memref.store %v0, %temp2[%c0, %c0, %c0] : memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
 
   // CHECK: %[[ALLOC1:.*]] = memref.alloc(%[[ARG0]])
   %3 = memref.alloc(%arg0) : memref<?x16x4xf32, offset : 0, strides : [64, 4, 1]>
   // Test: subview with constant operands but dynamic base memref is folded as long as the strides and offset of the base memref are static.
   // CHECK: memref.subview %[[ALLOC1]][0, 0, 0] [7, 11, 15] [1, 1, 1] :
   // CHECK-SAME: memref<?x16x4xf32, #[[$BASE_MAP0]]>
   // CHECK-SAME: to memref<7x11x15xf32, #[[$BASE_MAP0]]>
   %4 = memref.subview %3[%c0, %c0, %c0] [%c7, %c11, %c15] [%c1, %c1, %c1]
     : memref<?x16x4xf32, offset : 0, strides : [64, 4, 1]> to
       memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
-  memref.store %v0, %4[%c0, %c0, %c0] : memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
+  %temp4 = call @break_chain(%4) : (memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>) -> (memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>)
+  memref.store %v0, %temp4[%c0, %c0, %c0] : memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>
 
   // Test: subview offset operands are folded correctly w.r.t. base strides.
   // CHECK: memref.subview %[[ALLOC0]][1, 2, 7] [7, 11, 2] [1, 1, 1] :
   // CHECK-SAME: memref<8x16x4xf32, #[[$BASE_MAP0]]> to
   // CHECK-SAME: memref<7x11x2xf32, #[[$SUBVIEW_MAP1]]>
   %5 = memref.subview %0[%c1, %c2, %c7] [%c7, %c11, %c2] [%c1, %c1, %c1]
     : memref<8x16x4xf32, offset : 0, strides : [64, 4, 1]> to
       memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]>