[mlir][Standard] Canonicalize chains of tensor_cast operations

[herhut]

Summary

Adds a pattern that replaces a chain of two tensor_cast operations by a single
tensor_cast operation if doing so will not remove constraints on the shapes.

[herhut]

I have considered to place a templated version of joinShapes somewhere (to be used with ShapedType and a template parameter to figure out the type to create) but wasn't sure where and whether useful.

[jpienaar]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3143

This is a little restrictive, especially if one considers element types with nested types. Something operating on shapes instead could be reused and no need for templating or special restrictions on element type. E.g., join of two shapes with result being a vector/populating a vector (where first element is rank to allow returning unranked) or using ShapedComponentType (https://github.com/llvm/llvm-project/blob/8985755762a429573af2ce657274772339d3b9db/mlir/include/mlir/Interfaces/InferTypeOpInterface.h#L34). The latter does have an element type and attribute extra (the former can be left empty unless equal in the join, the latter we haven't used yet and so remove it too)

3197

Perhaps expand to say that it won't be used if it contains less information and remove empty line below.

[ftynse]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3150

Drop llvm::

3151

assert that ranks are equal

3162

Reserve space before pushing in a loop, you are guaranteed to push back exactly once on each iteration

[frgossen]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3142

I think, you could get away without materialising the joined shapes here.

If the cast chain is A -> B -> C then
A <= B or C <= B is a sufficient condition for this rewrite (smaller meaning more concrete shapes).

3150

Would resize(one.getRank()) be useful?

3200

A comment for why this is needed could help.
Cannot eliminate intermediate cast if is stricter than the resulting cast.

[herhut]

Comments

[herhut]

I will do the change to use ShapedTypeComponents in a follow up to enable reuse for the memref case.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3142

That is true, if computed for each element of the shape. It is not required for the shapes as a whole. Like A = [1,?], B = [?,2] and C = [1,2]. So implementing lessPrecise(ShapeType, ShapeType) and composing the above from it does not work. So I need a specialized castingAtoCisEquivalentToCastingAtoBtoC(), which I shied away from as join seemed more reusable.

So how much to we value performance over code reuse?

3143

A thanks! How about moving ShapedTypeComponents to TypeUtilities.h or even cook up a ShapeUtilities.h? If we use this more broadly than just the InferTypeOpInterface, it should live in a shared location.

[frgossen]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3142

Right, the order on shapes is not complete.

The join version is a lot easier to read :-)

[herhut]

This has landed a while back in 5e0ded268929b87ddf2c5e077c9185554342f602.

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

def TensorCastOp : CastOp<"tensor_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 tensor_cast is always a tensor.
     TensorType getType() { return getResult().getType().cast<TensorType>(); }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TensorLoadOp
 //===----------------------------------------------------------------------===//
 
 def TensorLoadOp : Std_Op<"tensor_load",
     [SameOperandsAndResultShape, SameOperandsAndResultElementType,

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

bool TensorCastOp::areCastCompatible(Type a, Type b) {
 
   return succeeded(verifyCompatibleShape(aT, bT));
 }
 
 OpFoldResult TensorCastOp::fold(ArrayRef<Attribute> operands) {
   return impl::foldCastOp(*this);
 }
 
+/// Compute a TensorType that has the joined shape knowledge of the two
+/// given TensorTypes. The element types need to match.
+static TensorType joinShapes(TensorType one, TensorType two) {
+  assert(one.getElementType() == two.getElementType());
+
+  if (!one.hasRank())
+    return two;
+  if (!two.hasRank())
+    return one;
+
+  int64_t rank = one.getRank();
+  if (rank != two.getRank())
+    return {};
+
+  SmallVector<int64_t, 4> join;
+  join.reserve(rank);
+  for (int64_t i = 0; i < rank; ++i) {
+    if (one.isDynamicDim(i)) {
+      join.push_back(two.getDimSize(i));
+      continue;
+    }
+    if (two.isDynamicDim(i)) {
+      join.push_back(one.getDimSize(i));
+      continue;
+    }
+    if (one.getDimSize(i) != two.getDimSize(i))
+      return {};
+    join.push_back(one.getDimSize(i));
+  }
+  return RankedTensorType::get(join, one.getElementType());
+}
+
+namespace {
+
+/// Replaces chains of two tensor_cast operations by a single tensor_cast
+/// operation if doing so does not remove runtime constraints.
+struct ChainedTensorCast : public OpRewritePattern<TensorCastOp> {
+  using OpRewritePattern<TensorCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TensorCastOp tensorCast,
+                                PatternRewriter &rewriter) const final {
+    auto tensorCastOperand =
+        tensorCast.getOperand().getDefiningOp<TensorCastOp>();
+
+    if (!tensorCastOperand)
+      return failure();
+
+    auto sourceType =
+        tensorCastOperand.getOperand().getType().cast<TensorType>();
+    auto intermediateType = tensorCastOperand.getType().cast<TensorType>();
+    auto resultType = tensorCast.getType().cast<TensorType>();
+
+    // We can remove the intermediate cast if joining all three produces the
+    // same result as just joining the source and result shapes.
+    auto firstJoin =
+        joinShapes(joinShapes(sourceType, intermediateType), resultType);
+
+    // The join might not exist if the cast sequence would fail at runtime.
+    if (!firstJoin)
+      return failure();
+
+    // The newJoin always exists if the above join exists, it might just contain
+    // less information. If so, we cannot drop the intermediate cast, as doing
+    // so would remove runtime checks.
+    auto newJoin = joinShapes(sourceType, resultType);
+    if (firstJoin != newJoin)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<TensorCastOp>(tensorCast, resultType,
+                                              tensorCastOperand.getOperand());
+    return success();
+  }
+};
+
+} // namespace
+
+void TensorCastOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<ChainedTensorCast>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // Helpers for Tensor[Load|Store]Op
 //===----------------------------------------------------------------------===//
 
 static Type getTensorTypeFromMemRefType(Type type) {
   if (auto memref = type.dyn_cast<MemRefType>())
     return RankedTensorType::get(memref.getShape(), memref.getElementType());
   if (auto memref = type.dyn_cast<UnrankedMemRefType>())

mlir/test/Transforms/canonicalize.mlir

%0 = dynamic_tensor_from_elements %size1, %c5, %size4 {
     %1 = constant 32 : index
     yield %1 : index
   // CHECK: : tensor<3x?x5x7x?xindex>
   } : tensor<3x?x?x7x?xindex>
   // CHECK: tensor_cast %{{.*}} : tensor<3x?x5x7x?xindex> to tensor<3x?x?x7x?xindex>
   return %0 : tensor<3x?x?x7x?xindex>
 }
 
+// -----
+
+// CHECK-LABEL: @tensor_cast_chain_ok
+// CHECK-SAME: %[[IN:.*]]: tensor<*xi32>
+func @tensor_cast_chain_ok(%input: tensor<*xi32>) -> tensor<4x8xi32> {
+  // CHECK-NEXT: %[[RES:.*]] = tensor_cast %[[IN]] : tensor<*xi32> to tensor<4x8xi32>
+  %0 = tensor_cast %input : tensor<*xi32> to tensor<4x?xi32>
+  %1 = tensor_cast %0 : tensor<4x?xi32> to tensor<4x8xi32>
+  // CHECK-NEXT: return %[[RES]]
+  return %1 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensor_cast_chain_regain
+// CHECK-SAME: %[[IN:.*]]: tensor<4xi32>
+func @tensor_cast_chain_regain(%input: tensor<4xi32>) -> tensor<4xi32> {
+  %0 = tensor_cast %input : tensor<4xi32> to tensor<?xi32>
+  %1 = tensor_cast %0 : tensor<?xi32> to tensor<4xi32>
+  // CHECK-NEXT: return %[[IN]]
+  return %1 : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensor_cast_chain_keep
+// CHECK-SAME: %[[IN:.*]]: tensor<?x?xi32>
+func @tensor_cast_chain_keep(%input: tensor<?x?xi32>) -> tensor<?x8xi32> {
+  // CHECK-NEXT: %[[C1:.*]] = tensor_cast %[[IN]]
+  %0 = tensor_cast %input : tensor<?x?xi32> to tensor<4x?xi32>
+  // CHECK-NEXT: %[[C2:.*]] = tensor_cast %[[C1]]
+  %1 = tensor_cast %0 : tensor<4x?xi32> to tensor<?x8xi32>
+  // CHECK-NEXT: return %[[C2]]
+  return %1 : tensor<?x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensor_cast_chain_invalid
+// CHECK-SAME: %[[IN:.*]]: tensor<4x8xi32>
+func @tensor_cast_chain_invalid(%input: tensor<4x8xi32>) -> tensor<8x4xi32> {
+  // CHECK-NEXT: %[[C1:.*]] = tensor_cast %[[IN]]
+  %0 = tensor_cast %input : tensor<4x8xi32> to tensor<?x?xi32>
+  // CHECK-NEXT: %[[C2:.*]] = tensor_cast %[[C1]]
+  %1 = tensor_cast %0 : tensor<?x?xi32> to tensor<8x4xi32>
+  // CHECK-NEXT: return %[[C2]]
+  return %1 : tensor<8x4xi32>
+}