[mlir][MemRef] Simplify extract_strided_metadata(reinterpret_cast)

[qcolombet]

Summary

This patch adds a pattern to simplify

base, offset, sizes, strides =
  extract_strided_metadata(
    reinterpret_cast(src, srcOffset, srcSizes, srcStrides))

Into

base, baseOffset, ... = extract_strided_metadata(src)
offset = srcOffset
sizes = srcSizes
strides = srcStrides

Note: Reinterpret_cast with unranked sources are not simplified since they cannot feed extract_strided_metadata operations.
Note2: This patch is part of a series that replaces the lowering of view like operations in MemRefToLLVM with composable patterns.
Note3: This patch depends on other patches in that series but only for the output of the tests.

[qcolombet]

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

687

Is there a generic way of doing that?
Something like ExtractStridedMetadata::verifySourceType?

[nicolasvasilache]

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

667

This feels very fishy .. I would expect just srcOffset here.

Computing the sum is a "subview"-like semantics, reinterpret_cast should just specify the offset UIAM

687

Implement your own ?

722

I am really not sure about this ..

[qcolombet]

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

667

That makes sense.
I thought it was adding say offset to the given src_memref, but indeed that's not really "reinterpret-y".

Fixing.

722

You're right, this is plain wrong :).

[qcolombet]

• Just forward the source offset. Reinterpret_cast doesn't do any math.

[qcolombet]

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

667

Fixed.

687

I'll look into that as a follow up patch if that works for you, since if this doesn't exist yet, it will likely require some tablegen work.

[qcolombet]

• Use inferReturnTypes to check if an extract_strided_metadata operation can be built instead of hardcoding the MemRefType.

[qcolombet]

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

687

Done that part by reusing the inferReturnTypes API that we enabled for these operations at https://reviews.llvm.org/D135734.

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

if (!viewLikeOp)
       return rewriter.notifyMatchFailure(extractOp, "not a ViewLike source");
     rewriter.updateRootInPlace(extractOp, [&]() {
       extractOp.getSourceMutable().assign(viewLikeOp.getViewSource());
     });
     return success();
   }
 };
 
+/// Replace `base, offset, sizes, strides =
+///              extract_strided_metadata(
+///                 reinterpret_cast(src, srcOffset, srcSizes, srcStrides))`
+/// With
+/// ```
+/// base, ... = extract_strided_metadata(src)
+/// offset = srcOffset
+/// sizes = srcSizes
+/// strides = srcStrides
+/// ```
+///
+/// In other words, consume the `reinterpret_cast` and apply its effects
+/// on the offset, sizes, and strides.
+class ExtractStridedMetadataOpReinterpretCastFolder
+    : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ExtractStridedMetadataOp extractStridedMetadataOp,
+                  PatternRewriter &rewriter) const override {
+    auto reinterpretCastOp = extractStridedMetadataOp.getSource()
+                                 .getDefiningOp<memref::ReinterpretCastOp>();
+    if (!reinterpretCastOp)
+      return failure();
+
+    Location loc = extractStridedMetadataOp.getLoc();
+    // Check if the source is suitable for extract_strided_metadata.
+    SmallVector<Type> inferredReturnTypes;
+    if (failed(extractStridedMetadataOp.inferReturnTypes(
+            rewriter.getContext(), loc, {reinterpretCastOp.getSource()},
+            /*attributes=*/{}, /*regions=*/{}, inferredReturnTypes)))
+      return rewriter.notifyMatchFailure(
+          reinterpretCastOp, "reinterpret_cast source's type is incompatible");
+
+    auto memrefType =
+        reinterpretCastOp.getResult().getType().cast<MemRefType>();
+    unsigned rank = memrefType.getRank();
+    SmallVector<OpFoldResult> results;
+    results.resize_for_overwrite(rank * 2 + 2);
+
+    auto newExtractStridedMetadata =
+        rewriter.create<memref::ExtractStridedMetadataOp>(
+            loc, reinterpretCastOp.getSource());
+
+    // Register the base_buffer.
+    results[0] = newExtractStridedMetadata.getBaseBuffer();
+
+    // Register the new offset.
+    results[1] = getValueOrCreateConstantIndexOp(
+        rewriter, loc, reinterpretCastOp.getMixedOffsets()[0]);
+
+    const unsigned sizeStartIdx = 2;
+    const unsigned strideStartIdx = sizeStartIdx + rank;
+
+    SmallVector<OpFoldResult> sizes = reinterpretCastOp.getMixedSizes();
+    SmallVector<OpFoldResult> strides = reinterpretCastOp.getMixedStrides();
+    for (unsigned i = 0; i < rank; ++i) {
+      results[sizeStartIdx + i] = sizes[i];
+      results[strideStartIdx + i] = strides[i];
+    }
+    rewriter.replaceOp(extractStridedMetadataOp,
+                       getValueOrCreateConstantIndexOp(rewriter, loc, results));
+    return success();
+  }
+};
+
 /// Replace `base, offset =
 ///            extract_strided_metadata(extract_strided_metadata(src)#0)`
 /// With
 /// ```
 /// base, ... = extract_strided_metadata(src)
 /// offset = 0
 /// ```
 class ExtractStridedMetadataOpExtractStridedMetadataFolder
patterns
       .add<ExtractStridedMetadataOpSubviewFolder,
            ExtractStridedMetadataOpReshapeFolder<
                memref::ExpandShapeOp, getExpandedSizes, getExpandedStrides>,
            ExtractStridedMetadataOpReshapeFolder<
                memref::CollapseShapeOp, getCollapsedSize, getCollapsedStride>,
            ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
            ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
            RewriteExtractAlignedPointerAsIndexOfViewLikeOp,
+           ExtractStridedMetadataOpReinterpretCastFolder,
            ExtractStridedMetadataOpExtractStridedMetadataFolder>(
           patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

%base, %offset =
     -> memref<i32>, index
   %base2, %offset2 =
     memref.extract_strided_metadata %base:memref<i32>
     -> memref<i32>, index
 
   return %base2, %offset2 :
       memref<i32>, index
 }
+
+// -----
+
+// Check that we simplify extract_strided_metadata of reinterpret_cast
+// when the source of the reinterpret_cast is compatible with what
+// `extract_strided_metadata`s accept.
+//
+// When we apply the transformation the resulting offset, sizes and strides
+// should come straight from the inputs of the reinterpret_cast.
+//
+// CHECK-LABEL: func @extract_strided_metadata_of_reinterpret_cast
+//  CHECK-SAME: %[[ARG:.*]]: memref<?x?xi32, strided<[?, ?], offset: ?>>, %[[DYN_OFFSET:.*]]: index, %[[DYN_SIZE0:.*]]: index, %[[DYN_SIZE1:.*]]: index, %[[DYN_STRIDE0:.*]]: index, %[[DYN_STRIDE1:.*]]: index)
+//
+//       CHECK: %[[BASE:.*]], %{{.*}}, %{{.*}}:2, %{{.*}}:2 = memref.extract_strided_metadata %[[ARG]]
+//
+//       CHECK: return %[[BASE]], %[[DYN_OFFSET]], %[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_STRIDE0]], %[[DYN_STRIDE1]]
+func.func @extract_strided_metadata_of_reinterpret_cast(
+  %arg : memref<?x?xi32, strided<[?, ?], offset:?>>,
+  %offset: index,
+  %size0 : index, %size1 : index,
+  %stride0 : index, %stride1 : index)
+  -> (memref<i32>, index,
+      index, index,
+      index, index) {
+
+  %cast =
+    memref.reinterpret_cast %arg to
+      offset: [%offset],
+      sizes: [%size0, %size1],
+      strides: [%stride0, %stride1] :
+      memref<?x?xi32, strided<[?, ?], offset: ?>> to
+      memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  %base, %base_offset, %sizes:2, %strides:2 =
+    memref.extract_strided_metadata %cast:memref<?x?xi32, strided<[?, ?], offset: ?>>
+    -> memref<i32>, index,
+       index, index,
+       index, index
+
+  return %base, %base_offset,
+    %sizes#0, %sizes#1,
+    %strides#0, %strides#1 :
+      memref<i32>, index,
+      index, index,
+      index, index
+}
+
+// -----
+
+// Check that we don't simplify extract_strided_metadata of
+// reinterpret_cast when the source of the cast is unranked.
+// Unranked memrefs cannot feed into extract_strided_metadata operations.
+// Note: Technically we could still fold the sizes and strides.
+//
+// CHECK-LABEL: func @extract_strided_metadata_of_reinterpret_cast_unranked
+//  CHECK-SAME: %[[ARG:.*]]: memref<*xi32>, %[[DYN_OFFSET:.*]]: index, %[[DYN_SIZE0:.*]]: index, %[[DYN_SIZE1:.*]]: index, %[[DYN_STRIDE0:.*]]: index, %[[DYN_STRIDE1:.*]]: index)
+//
+//       CHECK: %[[CAST:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [%[[DYN_OFFSET]]], sizes: [%[[DYN_SIZE0]], %[[DYN_SIZE1]]], strides: [%[[DYN_STRIDE0]], %[[DYN_STRIDE1]]]
+//       CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[CAST]]
+//
+//       CHECK: return %[[BASE]], %[[OFFSET]], %[[SIZES]]#0, %[[SIZES]]#1, %[[STRIDES]]#0, %[[STRIDES]]#1
+func.func @extract_strided_metadata_of_reinterpret_cast_unranked(
+  %arg : memref<*xi32>,
+  %offset: index,
+  %size0 : index, %size1 : index,
+  %stride0 : index, %stride1 : index)
+  -> (memref<i32>, index,
+      index, index,
+      index, index) {
+
+  %cast =
+    memref.reinterpret_cast %arg to
+      offset: [%offset],
+      sizes: [%size0, %size1],
+      strides: [%stride0, %stride1] :
+      memref<*xi32> to
+      memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  %base, %base_offset, %sizes:2, %strides:2 =
+    memref.extract_strided_metadata %cast:memref<?x?xi32, strided<[?, ?], offset: ?>>
+    -> memref<i32>, index,
+       index, index,
+       index, index
+
+  return %base, %base_offset,
+    %sizes#0, %sizes#1,
+    %strides#0, %strides#1 :
+      memref<i32>, index,
+      index, index,
+      index, index
+}
+
+// -----
+
+// Similar to @extract_strided_metadata_of_reinterpret_cast, just make sure
+// we handle 0-D properly.
+//
+// CHECK-LABEL: func @extract_strided_metadata_of_reinterpret_cast_rank0
+//  CHECK-SAME: %[[ARG:.*]]: memref<i32, strided<[], offset: ?>>, %[[DYN_OFFSET:.*]]: index, %[[DYN_SIZE0:.*]]: index, %[[DYN_SIZE1:.*]]: index, %[[DYN_STRIDE0:.*]]: index, %[[DYN_STRIDE1:.*]]: index)
+//
+//       CHECK: %[[BASE:.*]], %[[BASE_OFFSET:.*]] = memref.extract_strided_metadata %[[ARG]]
+//
+//       CHECK: return %[[BASE]], %[[DYN_OFFSET]], %[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_STRIDE0]], %[[DYN_STRIDE1]]
+func.func @extract_strided_metadata_of_reinterpret_cast_rank0(
+  %arg : memref<i32, strided<[], offset:?>>,
+  %offset: index,
+  %size0 : index, %size1 : index,
+  %stride0 : index, %stride1 : index)
+  -> (memref<i32>, index,
+      index, index,
+      index, index) {
+
+  %cast =
+    memref.reinterpret_cast %arg to
+      offset: [%offset],
+      sizes: [%size0, %size1],
+      strides: [%stride0, %stride1] :
+      memref<i32, strided<[], offset: ?>> to
+      memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  %base, %base_offset, %sizes:2, %strides:2 =
+    memref.extract_strided_metadata %cast:memref<?x?xi32, strided<[?, ?], offset: ?>>
+    -> memref<i32>, index,
+       index, index,
+       index, index
+
+  return %base, %base_offset,
+    %sizes#0, %sizes#1,
+    %strides#0, %strides#1 :
+      memref<i32>, index,
+      index, index,
+      index, index
+}