[mlir][MemRef] Canonicalize reinterpret_cast(extract_strided_metadata)

[qcolombet]

Summary

Add a canonicalizetion step for
reinterpret_cast(extract_strided_metadata).
This step replaces this sequence of operations by either:

• A noop, i.e., the original memref is directly used, or
• A plain cast of the original memref

The choice is ultimately made based on whether the original memref type
is equal to what the reinterpret_cast is producing. For instance, the
reinterpret_cast could be changing some dimensions from static to
dynamic and in such case, we need to keep a cast.

The transformation is currently only performed when the reinterpret_cast
uses exactly the same arguments as what the extract_strided_metadata
produces. It may be possible to be more aggressive here but I wanted to
start with a relatively simple MLIR patch for my first one!

[nicolasvasilache]

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

1622

We're now allowed to use C++17 auto, something like this should work here and below.

for (auto [extractStride, reinterpretStride] : llvm::zip(extractStridedMetadata.getStrides(), op.getStrides()))))
  if (extractStride != reinterpretStride)
    return failure();

[qcolombet]

Use C++17 auto style

[qcolombet]

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

1622

Nice!

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_RankOp : MemRef_Op<"rank", [NoSideEffect]> {
   let summary = "rank operation";

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

OpFoldResult ReinterpretCastOp::fold(ArrayRef<Attribute> /*operands*/) {
   if (auto prevSrc = getPrevSrc()) {
     getSourceMutable().assign(prevSrc);
     return getResult();
   }
 
   return nullptr;
 }
 
+namespace {
+/// Replace reinterpret_cast(extract_strided_metadata memref) -> memref.
+struct ReinterpretCastOpExtractStridedMetadataFolder
+    : public OpRewritePattern<ReinterpretCastOp> {
+public:
+  using OpRewritePattern<ReinterpretCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ReinterpretCastOp op,
+                                PatternRewriter &rewriter) const override {
+    auto extractStridedMetadata =
+        op.getSource().getDefiningOp<ExtractStridedMetadataOp>();
+    if (!extractStridedMetadata)
+      return failure();
+    // Check if the reinterpret cast reconstructs a memref with the exact same
+    // properties as the extract strided metadata.
+
+    // First, check that the strides are the same.
+    if (extractStridedMetadata.getStrides().size() != op.getStrides().size())
+      return failure();
+    for (auto [extractStride, reinterpretStride] :
+         llvm::zip(extractStridedMetadata.getStrides(), op.getStrides()))
+      if (extractStride != reinterpretStride)
+        return failure();
+
+    // Second, check the sizes.
+    if (extractStridedMetadata.getSizes().size() != op.getSizes().size())
+      return failure();
+    for (auto [extractSize, reinterpretSize] :
+         llvm::zip(extractStridedMetadata.getSizes(), op.getSizes()))
+      if (extractSize != reinterpretSize)
+        return failure();
+
+    // Finally, check the offset.
+    if (op.getOffsets().size() != 1 &&
+        extractStridedMetadata.getOffset() != *op.getOffsets().begin())
+      return failure();
+
+    // At this point, we know that the back and forth between extract strided
+    // metadata and reinterpret cast is a noop. However, the final type of the
+    // reinterpret cast may not be exactly the same as the original memref.
+    // E.g., it could be changing a dimension from static to dynamic. Check that
+    // here and add a cast if necessary.
+    Type srcTy = extractStridedMetadata.getSource().getType();
+    if (srcTy == op.getResult().getType())
+      rewriter.replaceOp(op, extractStridedMetadata.getSource());
+    else
+      rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(),
+                                          extractStridedMetadata.getSource());
+
+    return success();
+  }
+};
+} // namespace
+
+void ReinterpretCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                    MLIRContext *context) {
+  results.add<ReinterpretCastOpExtractStridedMetadataFolder>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // Reassociative reshape ops
 //===----------------------------------------------------------------------===//
 
 /// Helper function for verifying the shape of ExpandShapeOp and ResultShapeOp
 /// result and operand. Layout maps are verified separately.
 ///
 /// If `allowMultipleDynamicDimsPerGroup`, multiple dynamic dimensions are

mlir/test/Dialect/MemRef/canonicalize.mlir

 func.func @reinterpret_of_subview(%arg : memref<?xi8>, %size1: index, %size2: index) -> memref<?xi8> {
   %0 = memref.subview %arg[0] [%size1] [1] : memref<?xi8> to memref<?xi8>
   %1 = memref.reinterpret_cast %0 to offset: [0], sizes: [%size2], strides: [1] : memref<?xi8> to memref<?xi8>
   return %1 : memref<?xi8>
 }
 
 // -----
 
+// Check that a reinterpret cast of an equivalent extract strided metadata
+// is canonicalized to a plain cast when the destination type is different
+// than the type of the original memref.
+// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_type_mistach
+//  CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
+//       CHECK: %[[CAST:.*]] = memref.cast %[[ARG]] : memref<8x2xf32> to memref<?x?xf32,
+//       CHECK: return %[[CAST]]
+func.func @reinterpret_of_extract_strided_metadata_w_type_mistach(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, offset : ?, strides : [?, ?]> {
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
+  %m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+  return %m2 : memref<?x?xf32, offset: ?, strides: [?, ?]>
+}
+
+// -----
+
+// Check that a reinterpret cast of an equivalent extract strided metadata
+// is completely removed when the original memref has the same type.
+// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_same_type
+//  CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
+//       CHECK: return %[[ARG]]
+func.func @reinterpret_of_extract_strided_metadata_same_type(%arg0 : memref<8x2xf32>) -> memref<8x2xf32> {
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
+  %m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<8x2xf32>
+  return %m2 : memref<8x2xf32>
+}
+
+// -----
+
+// Check that we don't simplify reinterpret cast of extract strided metadata
+// when the strides don't match.
+// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_different_stride
+//  CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
+//       CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]]
+//       CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[OFFSET]]], sizes: [4, 2, 2], strides: [1, 1, %[[STRIDES]]#1]
+//       CHECK: return %[[RES]]
+func.func @reinterpret_of_extract_strided_metadata_w_different_stride(%arg0 : memref<8x2xf32>) -> memref<?x?x?xf32, offset : ?, strides : [?, ?, ?]> {
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
+  %m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [4, 2, 2], strides: [1, 1, %strides#1] : memref<f32> to memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]>
+  return %m2 : memref<?x?x?xf32, offset: ?, strides: [?, ?, ?]>
+}
+// -----
+
+// Check that we don't simplify reinterpret cast of extract strided metadata
+// when the offset doesn't match.
+// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_different_offset
+//  CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
+//       CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]]
+//       CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [1], sizes: [%[[SIZES]]#0, %[[SIZES]]#1], strides: [%[[STRIDES]]#0, %[[STRIDES]]#1]
+//       CHECK: return %[[RES]]
+func.func @reinterpret_of_extract_strided_metadata_w_different_offset(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, offset : ?, strides : [?, ?]> {
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
+  %m2 = memref.reinterpret_cast %base to offset: [1], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+  return %m2 : memref<?x?xf32, offset: ?, strides: [?, ?]>
+}
+
+// -----
+
 func.func @canonicalize_rank_reduced_subview(%arg0 : memref<8x?xf32>,
     %arg1 : index) -> memref<?xf32, offset : ?, strides : [?]> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %0 = memref.subview %arg0[%c0, %c0] [1, %arg1] [%c1, %c1] : memref<8x?xf32> to memref<?xf32, offset : ?, strides : [?]>
   return %0 :  memref<?xf32, offset : ?, strides : [?]>
 }
 //  CHECK-DAG: #[[MAP:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
 //      CHECK: func @canonicalize_rank_reduced_subview
 // CHECK-SAME:     %[[ARG0:.+]]: memref<8x?xf32>
 // CHECK-SAME:     %[[ARG1:.+]]: index
 //      CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, 0] [1, %[[ARG1]]] [1, 1]
 // CHECK-SAME:       memref<8x?xf32> to memref<?xf32, #[[MAP]]>