[mlir] Add a canonicalization pattern for MemRefCastOp into dynamic MemRefs

[nicolasvasilache]

Summary

This diff adds a canonicalization pattern to canonicalize either:

mlir
  ... = memref_cast ... : memref<8x16xf32> to memref<?x?xf32>
  ... = memref_cast ... : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
        to memref<?x?xf32>

into

mlir
  ... = memref_cast ... : ... to memref<8x16xf32>

This diff additionally needs to modify a seemingly unrelated things that does
not play nicely atm. Some tests that return memref need to be changed because
canonicalization does not propagate to function signature.

[nicolasvasilache]

Remove spurious declarations.

[rriddle]

mlir/lib/Dialect/StandardOps/Ops.cpp

1804

You are replacing with something of a different type. That doesn't really fit as a canonicalization. For this to be conservatively correct, it would mean you need to prevent a memref cast from ever being used: in a function call, as a branch argument, etc.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[nmostafa]

mlir/lib/Dialect/StandardOps/Ops.cpp

1768

You need to handle or bailout on UnrankedMemRefType (and probably add a test for it). Also, the type propagation is only valid based on how the result is used : e.g. if a func argument, then the conversion is invalid.

[nicolasvasilache]

Address review comments.

[nicolasvasilache]

Add comment.

[nicolasvasilache]

mlir/lib/Dialect/StandardOps/Ops.cpp

1768

Thanks, I addressed those points but there are quite deeper implications as River pointed out.

1804

How fundamental do you foresee this to be a limitation in the longer term?
The use cases look like:

cast -> X1 
...
cast -> Xk

where k is small(ish).
Very concretely, this enables n-D linag.vectorization to kick-in as a declarative pattern rewrite that composes nicely with all the other things.
So at some point I will want a way to hook that up with Ops that should not be seen here.

I could extend this into a conditional pattern that is called with a lambda from the proper places.
It also seems OpInterfaces would be a natural fit but it seems a very unnatural thing to say "set of ops for which memrefcast can canonicalize its return type".

OTOH I can also bail here and make this a rewrite pattern locally in the pass that I am experimenting with.
But the way things start looking like passes become a union of loose connected patterns that compose nicely and locally to achieve global effects.
The most natural thing that this pattern made me think of is a canonicalization (modulo the type change), which is why I am trying this maybe unnatural thing.

Thoughts?

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 61909 tests passed, 0 failed and 782 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[nicolasvasilache]

In addition to all this I seem to have also uncovered issues in the existing memref_cast folders .. investigating.
It is likely the fact that memref_cast and other casts are considered to be "almost the same thing" is completely wrong.

[nicolasvasilache]

Scratch that, I am wrong, sorry .. the load(memref) just forwards the operand and that is ok.
It seems my current pattern is not properly testable without the ops I want.

I will move it into Linalg Transforms for now and get the end-to-end story more visible.
When/if it makes sense we can discuss whether canonicalization could ever make sense for this.

mlir/include/mlir/Dialect/StandardOps/Ops.td

def MemRefCastOp : CastOp<"memref_cast"> {
   let extraClassDeclaration = [{
     /// Return true if `a` and `b` are valid operand and result pairs for
     /// the operation.
     static bool areCastCompatible(Type a, Type b);
 
     /// The result of a memref_cast is always a memref.
     Type getType() { return getResult().getType(); }
   }];
+  
+  let hasCanonicalizer = 1;
 }
 
 def MulFOp : FloatArithmeticOp<"mulf"> {
   let summary = "floating point multiplication operation";
   let hasFolder = 1;
 }
 
 def MulIOp : IntArithmeticOp<"muli", [Commutative]> {

mlir/lib/Dialect/StandardOps/Ops.cpp

if (aMemSpace != bMemSpace)
       return false;
 
     return true;
   }
 
   return false;
 }
 
+namespace {
+
+/// Canonicalize either:
+///
+/// ```mlir
+///   ... = memref_cast ... : memref<8x16xf32> to memref<?x?xf32>
+///   ... = memref_cast ... : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
+//          to memref<?x?xf32>
+/// ```
+///
+/// into
+///
+/// ```mlir
+///   ... = memref_cast ... : ... to memref<8x16xf32>
+/// ```
+///
+/// One big caveat is that this canonicalization changes the result type and
+/// cannot work in a generic fashion at the moment. As a consequence, this
+/// inspects uses and bails out if it hits an op it does not know about.
+// TODO(ntv, rriddle): this may be a good use case for an OpInterface.
+class MemRefCastToDynamicFolder final : public OpRewritePattern<MemRefCastOp> {
+public:
+  using OpRewritePattern<MemRefCastOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(MemRefCastOp castOp,
+                                     PatternRewriter &rewriter) const override {
+    MemRefType sourceType = castOp.source().getType().dyn_cast<MemRefType>();
+    AffineMap sourceMap = (sourceType.getAffineMaps().empty())
+                              ? AffineMap()
+                              : sourceType.getAffineMaps().front();
+    MemRefType resultType = castOp.getType().dyn_cast<MemRefType>();
+    MemRefType desiredResultType =
+        MemRefType::get(sourceType.getShape(), sourceType.getElementType(), {});
+
+    // If we don't have MemRefType as source and destination, bail out.
+    if (!sourceType || !resultType)
+      return matchFailure();
+
+    // If we're already in canonical form all is good.
+    if (resultType == desiredResultType)
+      return matchFailure();
+
+    // If resultType has a map, it needs to be the same as the source type to
+    // canonicalize.
+    if (!resultType.getAffineMaps().empty() && sourceType != resultType)
+      return matchFailure();
+
+    // Inspect uses and bails out if it hits an op it does not know about.
+    // TODO(ntv, rriddle): this may be a good use case for an OpInterface.
+    for (auto &u : castOp.getResult().getUses()) {
+      if (isa<LoadOp>(u.getOwner()) || isa<StoreOp>(u.getOwner()))
+        continue;
+      return matchFailure();
+    }
+
+    // Ensure that:
+    //   1. source is static
+    //   2. source and target have the same rank (will be extended when needed)
+    //   3. if result is partially static, ensure sizes match.
+    if (!sourceType.hasStaticShape() ||
+        sourceType.getRank() != resultType.getRank())
+      return matchFailure();
+    for (auto it : llvm::zip(sourceType.getShape(), resultType.getShape())) {
+      auto sourceSize = std::get<0>(it);
+      auto resultSize = std::get<1>(it);
+      if (ShapedType::isDynamic(resultSize))
+        continue;
+      if (sourceSize != resultSize)
+        return matchFailure();
+    }
+
+    if (sourceMap) {
+      int64_t offset;
+      SmallVector<int64_t, 4> strides;
+      auto res = getStridesAndOffset(sourceType, strides, offset);
+      (void)res;
+      assert(succeeded(res));
+      auto stridedMap =
+          makeStridedLinearLayoutMap(strides, offset, sourceMap.getContext());
+      if (sourceMap != stridedMap)
+        return matchFailure();
+    }
+
+    rewriter.replaceOpWithNewOp<MemRefCastOp>(castOp, castOp.source(),
+                                              desiredResultType);
+
+    return matchSuccess();
+  }
+};
+
+} // end anonymous namespace
+
+void MemRefCastOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<MemRefCastToDynamicFolder>(context);
+}
+
 OpFoldResult MemRefCastOp::fold(ArrayRef<Attribute> operands) {
   return impl::foldCastOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Transforms/canonicalize.mlir

func @subview(%arg0 : index, %arg1 : index) -> (index, index) {
   %7 = dim %4, 0 : memref<?x?x?xf32,
        affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + d1 * s2 + d2 * s3 + s0)>>
   %8 = dim %4, 1 : memref<?x?x?xf32,
        affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + d1 * s2 + d2 * s3 + s0)>>
 
   // CHECK: return %[[C7]], %[[C11]]
   return %7, %8 : index, index
 }
+
+// CHECK-LABEL: func @memref_cast_canonicalize
+func @memref_cast_canonicalize(%arg0: memref<4x8xf32>) -> f32 {
+  //      CHECK: constant 0 : index
+  %c0 = constant 0 : index
+  // Goes away
+  %0 = memref_cast %arg0 : memref<4x8xf32> to memref<?x?xf32>
+
+  // CHECK-NEXT: load {{.*}} : memref<4x8xf32>
+  %1 = load %0[%c0, %c0] : memref<?x?xf32>
+  return %1 : f32
+}