[mlir][MemRef] Add folder for no-op `memref.subview`.

[mravishankar]

Summary

When offsets, sizes and strides are empty, the memref.subview can be
folded into the source. This cannot be done as a folder since a
memref.cast might be required to address the types. Add a
canonicalization pattern for this.

[mravishankar]

Add some comments.

[bondhugula]

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

635

This change is not covered by the commit summary and title.

[bondhugula]

The revision is making an undocumented change to the DimOp's folder.

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

635

The whole block above is also missing a comment completely as well.

1937

This isn't the right or efficient way to replace an op's result with another Value in a rewrite pattern. RAUW shouldn't be used here; instead:

rewriter.replaceOp(subViewOp, source);

[bondhugula]

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

1925

Consider renaming this to EraseTrivialSubView or TrivialSubViewFolder to be consistent with other such simplifiers.

[mravishankar]

Rebase and address comments.

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

1937

Good catch! Thanks!

[bondhugula]

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

1968–1969

Sorted order here.

1981

Typo: empty

1981

I didn't understand this part - it could have dynamic strides, offsets, and sizes, right?

[mravishankar]

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

1981

The size of static_offsets acounts for dynamic offsets as well. A -1 is added to static_offsets when the offset is dynamic. (same is true for sizes and strides). The semantics of the subview operation is that missing offsets implies offset is 0, missing sizes implies sizes is same as source size and missing strides implies strides are 1. So all offsets, sizes, and strides missing means that the subview is a no-op.

[mravishankar]

Address comments and rebase.

[mravishankar]

Fix typo

[mravishankar]

Rebase.

[mravishankar]

Rebase and fix test error.

[mravishankar]

FYI, closing this patch as not committing. Ideally memref.subview should not allow number of sizes, offsets and strides that are less than the rank of the source. This should be fixed. I am going to try to tackle that instead.

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

for (auto i : llvm::seq<unsigned>(0, sourceRank)) {
       if (unusedDims.count(i))
         continue;
       if (resultIndex == unsignedIndex) {
         sourceIndex = i;
         break;
       }
       resultIndex++;
     }
+    // If the source index is greater that the number of static sizes specified,
+    // then the semantics of the `SubViewOp` is that the size is same as the
+    // size of the the source. Do nothing here, and rely on subview op folder to
+    // fold away the subview itself.
+    if (sourceIndex >= subview.static_sizes().size())
+      return {};
     assert(subview.isDynamicSize(sourceIndex) &&
            "expected dynamic subview size");
     return subview.getDynamicSize(sourceIndex);
   }
 
   if (auto sizeInterface =
           dyn_cast_or_null<OffsetSizeAndStrideOpInterface>(definingOp)) {
     assert(sizeInterface.isDynamicSize(unsignedIndex) &&
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();
   }
 };
+
+/// Fold `memref.subview` with empty offsets, sizes and strides to the `source`
+/// (with casts to address type consistency).
+class TrivialSubViewOpFolder final : public OpRewritePattern<SubViewOp> {
+public:
+  using OpRewritePattern<SubViewOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SubViewOp subViewOp,
+                                PatternRewriter &rewriter) const override {
+    if (!subViewOp.static_offsets().empty() ||
+        !subViewOp.static_sizes().empty() ||
+        !subViewOp.static_strides().empty())
+      return failure();
+    Value source = subViewOp.source();
+    if (subViewOp.getSourceType() == subViewOp.getType()) {
+      rewriter.replaceOp(subViewOp, source);
+      return success();
+    }
+    rewriter.replaceOpWithNewOp<CastOp>(subViewOp, subViewOp.getType(), source);
+    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(), 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());
   }
 };
 
 void SubViewOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results
       .add<OpWithOffsetSizesAndStridesConstantArgumentFolder<
                SubViewOp, SubViewReturnTypeCanonicalizer, SubViewCanonicalizer>,
-           SubViewOpMemRefCastFolder>(context);
+           SubViewOpMemRefCastFolder, TrivialSubViewOpFolder>(context);
 }
 
 OpFoldResult SubViewOp::fold(ArrayRef<Attribute> operands) {
   auto resultShapedType = getResult().getType().cast<ShapedType>();
   auto sourceShapedType = source().getType().cast<ShapedType>();
 
   if (resultShapedType.hasStaticShape() &&
       resultShapedType == sourceShapedType) {
     return getViewSource();
   }
+  // If the offsets, sizes and strides are empty, the subview is a no-op.
+  if (static_offsets().empty() && static_sizes().empty() &&
+      static_strides().empty())
+    return getViewSource();
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/MemRef/canonicalize.mlir

%1 = memref.cast %0
         memref<1x4x1xf32, affine_map<(d0, d1, d2)[s0] -> (d0 * 35 + s0 + d1 * 7 + d2)>>
   return %1 : memref<1x4x1xf32, affine_map<(d0, d1, d2)[s0] -> (d0 * 35 + s0 + d1 * 7 + d2)>>
 }
 
 // CHECK-LABEL: func @reduced_memref
 //       CHECK:   %[[RESULT:.+]] = memref.subview
 //  CHECK-SAME:       memref<2x5x7x1xf32> to memref<1x4x1xf32, #{{.+}}>
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+func @no_op_scalar_subview(%arg0 : memref<f32>) -> memref<f32, affine_map<() -> (0)>> {
+  %0 = memref.subview %arg0[] [] [] : memref<f32> to memref<f32, affine_map<() -> (0)>>
+  return %0 : memref<f32, affine_map<() -> (0)>>
+}
+//      CHECK: func @no_op_scalar_subview(
+// CHECK-SAME:     %[[ARG0:.+]]: memref<f32>
+//      CHECK:   %[[CAST:.+]] = memref.cast %[[ARG0]]
+//      CHECK:   return %[[CAST]]

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
+
+// -----
+
+#map = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
+func @no_op_subview(%arg0 : memref<?x?xf32, #map>) -> index {
+  %c1 = arith.constant 1 : index
+  %0 = memref.subview %arg0[] [] [] : memref<?x?xf32, #map> to memref<?x?xf32, #map>
+  %1 = memref.dim %0, %c1 : memref<?x?xf32, #map>
+  return %1 : index
+}
+//      CHECK: func @no_op_subview(
+// CHECK-SAME:     %[[ARG0:.+]]: memref<?x?xf32
+//      CHECK:   %[[C1:.+]] = arith.constant 1 : index
+//      CHECK:   %[[RETURN:.+]] = memref.dim %[[ARG0]], %[[C1]]
+//      CHECK:   return %[[RETURN]]