[mlir][MemRef] Make reinterpret_cast(extract_strided_metadata) more robust

[qcolombet]

Summary

Prior to this patch the canonicalization pattern that turns reinterpret_cast(extract_strided_metadata) into cast was only applied when all the input operands of the reinterpret_cast are exactly all the output results of the extract_strided_metadata.

This missed simplification opportunities when the values would have hold the same constant values, but yet, come from different actual values.

E.g., prior to this patch, a pattern of the form:

%base, %offset = extract_strided_metadata %source : memref<i16>
reinterpret_cast %base to offset:[0]

Wouldn't have been simplified into a simple cast, because %offset is not directly the same value object as 0.

This patch teaches this pattern how to check if the constant values match what the results of the extract_strided_metadata operation would have hold.

Note: This patch will be helpful to keep nice MLIR outputs when I push the patch to do the forwarding of statically known information as part of a folding of extract_strided_metadata.

[nicolasvasilache]

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

1762

This feels like a very clunky API to use ..

Can we merge the OFR and int64_t into a single vector keeping the most static information and only directly compare equality ?

I can't tell offhand if whether this is possible or whether the logic allows for more complex mixings of Value vs int comparisons.

1773–1809

More doc plz, this is much more complex behavior now.

Basically the success cases are:

1. the source buffer of reinterpret_cast comes from extract_strided_metadata memref and
2. all offsets, sizes and strides are statically known to be equal.

Some small IR examples would be good too (fold case and cast case)

1850

multi-line code requires braces

[qcolombet]

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

1762

Can we merge the OFR and int64_t into a single vector keeping the most static information and only directly compare equality ?

I don't know how much it would help API-wise in the sense that the tricky bit is to get this information to beginning with. (I.e., what is currently done in sameValueOrConstant and more particularly the lambda getConstant.)

[qcolombet]

• Add comments

[qcolombet]

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

1762

Looking at this again, I don't see a whole lot of simplification without some code duplication (which is what I wanted to avoid in the first place).

The annoying part is, to get the most static information, we need both the OpFoldResult (to get the index value or to crack open the related arith.constant) and the related int64_t that we deduced from the type. I.e., in terms of API, we need both A and constantA to get the most static information. (And we also need to know if constant is dynamic.)

If we were to introduce a getMostStaticInformation it would essentially look like the getConstant lambda. Then for the checks themselves, we would have to replicate the logic that compares two "most static information". (I.e. what lives in sameValueOrConstant right now.

Ideas?

[nicolasvasilache]

Discussed offline, the way to go here is to add helpers to each op that encode the best efforts mixed static/dynamic semantics and are op specific.

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

1761

if we use zip here, this function can go away

1843

can we add this to the op itself to return a vector of opfoldresult which best effort contains the Value or Attr extract from the type.

1851

can we use a zip to drop the need for a vector in this complex API ?
Here and below

[qcolombet]

• Introduce getConstifiedMixedXXX methods for ReinterpretCastOp and ExtractStridedMetadataOp
• Use this new method to simplify the implementation of reinterpret_cast(extract_strided_metadata)

[qcolombet]

Hi @nicolasvasilache,

Please take another look.

I've introduced new methods to get the most constant values of the sizes, strides, and offset for extract_strided_metadata and reinterpret_cast (see getConstifiedMixedXXX methods; BTW suggestions for a better name are appreciated!)

The "complexity" is now all hidden in constifyIndexValues and the reinterpret_cast(extract_strided_metadata) combine looks straight forward (no zip, no function ptr indirection etc.).

Cheers,
-Quentin

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

930 ↗

(On Diff #472467)

We'll probably want these methods to be part of an interface.
ViewLikeInterface could have been a candidate but extract_strided_metadata doesn't conform to it, so we likely will need to introduce something else.
For now, I just added the methods where I needed them.

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

145

I acknowledge that this API is complicated, but it is not meant to be used directly.
This represents the meat of what the getConstifiedMixedXXX methods do, thus we have to support size, strides, and offset.

169

I'll do an RFC on how to properly fix that.
I have a wip patch for that, but it's getting longer than expected to fix all the places.

I wanted to unblock this patch, hence the workaround.
Let me know if you want me to file a bug for the proper fix.

1286

@nicolasvasilache
This refactoring is not required but allows to remove some code.
I'm happy to leave it out if you prefer the previous version.

1843

Done in getConstifiedMixedXXX methods.

[nicolasvasilache]

Sorry for the delay . let's go with this.
The API is indeed more complex than we'd like but at least it is well isolated form users so I am fine for now.

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

return op.emitOpError("dimension operand count does not equal memref "
                           "dynamic dimension count");
 
   unsigned numSymbols = 0;
   if (!memRefType.getLayout().isIdentity())
     numSymbols = memRefType.getLayout().getAffineMap().getNumSymbols();
   if (op.getSymbolOperands().size() != numSymbols)
     return op.emitOpError("symbol operand count does not equal memref symbol "
                           "count: expected ")
            << numSymbols << ", got " << op.getSymbolOperands().size();
 
   return success();
 }
 
 LogicalResult AllocOp::verify() { return verifyAllocLikeOp(*this); }
 
 LogicalResult AllocaOp::verify() {
   // An alloca op needs to have an ancestor with an allocation scope trait.
   if (!(*this)->getParentWithTrait<OpTrait::AutomaticAllocationScope>())
     return emitOpError(
         "requires an ancestor op with AutomaticAllocationScope trait");
 
   return verifyAllocLikeOp(*this);
 }
 
 namespace {
 /// Fold constant dimensions into an alloc like operation.
 template <typename AllocLikeOp>
 struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
   using OpRewritePattern<AllocLikeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AllocLikeOp alloc,
                                 PatternRewriter &rewriter) const override {
     // Check to see if any dimensions operands are constants.  If so, we can
     // substitute and drop them.
     if (llvm::none_of(alloc.getDynamicSizes(), [](Value operand) {
           return matchPattern(operand, matchConstantIndex());
         }))
       return failure();
 
     auto memrefType = alloc.getType();
 
if (!getSizes().empty()) {
     setNameFn(getSizes().front(), "sizes");
     setNameFn(getStrides().front(), "strides");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // GenericAtomicRMWOp
 //===----------------------------------------------------------------------===//
 
 void GenericAtomicRMWOp::build(OpBuilder &builder, OperationState &result,
                                Value memref, ValueRange ivs) {
   result.addOperands(memref);
   result.addOperands(ivs);
 
   if (auto memrefType = memref.getType().dyn_cast<MemRefType>()) {
     Type elementType = memrefType.getElementType();
     result.addTypes(elementType);
OpFoldResult ReinterpretCastOp::fold(ArrayRef<Attribute> /*operands*/) {
   if (auto prevSrc = getPrevSrc()) {
     getSourceMutable().assign(prevSrc);
     return getResult();
   }
 
   return nullptr;
 }
 
+/// Given a value \p A and \p B and their potential alternative constant
+/// representation \p constantA and \p constantB, check if \p A is equal to
+/// \p B.
+/// \p isDynamic is used to check if \p constantA and \p constantB are valid
+/// constants.
+static bool sameValueOrConstant(OpFoldResult A, int64_t constantA,
+                                OpFoldResult B, int64_t constantB,
+                                llvm::function_ref<bool(int64_t)> isDynamic) {
+  Value valA = A.is<Value>() ? A.get<Value>() : Value();
+  Value valB = B.is<Value>() ? B.get<Value>() : Value();
+  if (valA != Value() && valA == valB)
+    return true;
+
+  auto getConstant = [isDynamic](OpFoldResult val,
+                                 int64_t constant) -> Optional<int64_t> {
+    Optional<int64_t> maybeConstant;
+    if (isDynamic(constant)) {
+      if (val.is<Attribute>())
+        maybeConstant = val.get<Attribute>().cast<IntegerAttr>().getInt();
+      else
+        maybeConstant = getConstantIntValue(val.get<Value>());
+    } else {
+      maybeConstant = constant;
+    }
+    return maybeConstant;
+  };
+
+  Optional<int64_t> maybeConstantA = getConstant(A, constantA);
+  Optional<int64_t> maybeConstantB = getConstant(B, constantB);
+  if (!maybeConstantA || !maybeConstantB)
+    return false;
+  return *maybeConstantA == *maybeConstantB;
+}
+
+/// Check if the \p idx-th OpFoldResult of A and B are equal.
+/// \see sameValueOrConstant for more details.
+static bool sameValueOrConstantAtIdx(
+    ArrayRef<OpFoldResult> valuesA, ArrayRef<int64_t> constantsA,
+    ArrayRef<OpFoldResult> valuesB, ArrayRef<int64_t> constantsB,
+    llvm::function_ref<bool(int64_t)> isDynamic, unsigned idx) {
+  OpFoldResult A = valuesA[idx];
+  OpFoldResult B = valuesB[idx];
+  int64_t constantA = constantsA[idx];
+  int64_t constantB = constantsB[idx];
+  return sameValueOrConstant(A, constantA, B, constantB, isDynamic);
+}
+
 namespace {
-/// Replace reinterpret_cast(extract_strided_metadata memref) -> memref.
+/// Replace the sequence:
+/// ```
+/// base, offset, sizes, strides = extract_strided_metadata src
+/// dst = reinterpret_cast base to offset, sizes, strides
+/// ```
+/// With
+///
+/// ```
+/// dst = memref.cast src
+/// ```
+///
+/// Note: The cast operation is only inserted when the type of dst and src
+/// are not the same. E.g., when going from <4xf32> to <?xf32>.
+///
+/// This pattern also matches when the offset, sizes, and strides don't come
+/// directly from the `extract_strided_metadata`'s results but it can be
+/// statically proven that they would hold the same values.
+///
+/// For instance, the following sequence would be replaced:
+/// ```
+/// base, offset, sizes, strides =
+///   extract_strided_metadata memref : memref<3x4xty>
+/// dst = reinterpret_cast base to 0, [3, 4], strides
+/// ```
+/// Because we know (thanks to the type of the input memref) that variable
+/// `offset` and `sizes` will respectively hold 0 and [3, 4].
+///
+/// Similarly, the following sequence would be replaced:
+/// ```
+/// c0 = arith.constant 0
+/// c4 = arith.constant 4
+/// base, offset, sizes, strides =
+///   extract_strided_metadata memref : memref<3x4xty>
+/// dst = reinterpret_cast base to c0, [3, c4], strides
+/// ```
+/// Because we know that `offset`and `c0` will hold 0
+/// and `c4` will hold 4.
 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)
+    SmallVector<OpFoldResult, 4> extractStridesOfr =
+        getAsOpFoldResult(extractStridedMetadata.getStrides());
+    SmallVector<OpFoldResult, 4> reinterpretStridesOfr = op.getMixedStrides();
+    if (extractStridesOfr.size() != reinterpretStridesOfr.size())
+      return failure();
+
+    SmallVector<int64_t> reinterpretStrides;
+    int64_t reinterpretOffset;
+    bool hasKnownStridesAndOffset = succeeded(getStridesAndOffset(
+        op.getType(), reinterpretStrides, reinterpretOffset));
+    (void)hasKnownStridesAndOffset;
+    assert(hasKnownStridesAndOffset &&
+           "getStridesAndOffset must work on valid reinterpret_cast");
+    SmallVector<int64_t> extractStrides;
+    int64_t extractOffset;
+    auto sourceType =
+        extractStridedMetadata.getSource().getType().cast<MemRefType>();
+    hasKnownStridesAndOffset = succeeded(
+        getStridesAndOffset(extractStridedMetadata.getSource().getType(),
+                            extractStrides, extractOffset));
+    (void)hasKnownStridesAndOffset;
+    assert(hasKnownStridesAndOffset && "getStridesAndOffset must work on valid "
+                                       "extract_strided_metadata source");
+
+    unsigned rank = op.getType().getRank();
+    for (unsigned i = 0; i < rank; ++i) {
+      if (!sameValueOrConstantAtIdx(extractStridesOfr, extractStrides,
+                                    reinterpretStridesOfr, reinterpretStrides,
+                                    ShapedType::isDynamicStrideOrOffset, i))
         return failure();
+    }
 
     // Second, check the sizes.
-    if (extractStridedMetadata.getSizes().size() != op.getSizes().size())
+    assert(extractStridedMetadata.getSizes().size() ==
+               op.getMixedSizes().size() &&
+           "Strides and sizes rank must match");
+    ArrayRef<int64_t> reinterpretSizes = op.getType().getShape();
+    ArrayRef<int64_t> extractSizes = sourceType.getShape();
+    SmallVector<OpFoldResult, 4> extractSizesOfr =
+        getAsOpFoldResult(extractStridedMetadata.getSizes());
+    SmallVector<OpFoldResult, 4> reinterpretSizesOfr = op.getMixedSizes();
+    for (unsigned i = 0; i < rank; ++i) {
+      if (!sameValueOrConstantAtIdx(extractSizesOfr, extractSizes,
+                                    reinterpretSizesOfr, reinterpretSizes,
+                                    ShapedType::isDynamic, i))
         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())
+    assert(op.getMixedOffsets().size() == 1 &&
+           "reinterpret_cast with more than one offset should have been "
+           "rejected by the verifier");
+    if (!sameValueOrConstant(
+            getAsOpFoldResult(extractStridedMetadata.getOffset()),
+            extractOffset, op.getMixedOffsets()[0], reinterpretOffset,
+            ShapedType::isDynamicStrideOrOffset))
       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();

mlir/test/Dialect/MemRef/canonicalize.mlir

 func.func @reinterpret_of_extract_strided_metadata_w_type_mistach(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
   %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, strided<[?, ?], offset: ?>>
   return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
 }
 
 // -----
 
+// Similar to reinterpret_of_extract_strided_metadata_w_type_mistach except that
+// we check that the match happen when the static information has been folded.
+// E.g., in this case, we know that size of dim 0 is 8 and size of dim 1 is 2.
+// So even if we don't use the values sizes#0, sizes#1, as long as they have the
+// same constant value, the match is valid.
+// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_constants
+//  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_constants(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
+  %c8 = arith.constant 8: index
+  %m2 = memref.reinterpret_cast %base to offset: [0], sizes: [%c8, 2], strides: [2, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
+  return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
+}
+// -----
+
 // 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>