[mlir][Shape] Make sure tensor_cast(constant_shape) folding uses the correct type

[bkramer]

Summary

This is still subtle, but I think the test cases are sufficient to show
that it works.

[herhut]

This seems to do the right thing, I just don't understand why.

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

38

I am somewhat puzzled what this means. It creates a new Shape_ConstShapeOp with the given $ty which is the value? Why is it called ty then? How does the type come into this? Is it somehow inherited from the original operation?

[bkramer]

• Rename argument from $ty to $arg

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

38

$ty is just me coming up with a weird name. Renamed it to be more consistent.

The type is inherited from the original.

[silvas]

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

37

Maybe add a small comment here about how the types are preserved.

mlir/test/Dialect/Shape/canonicalize.mlir

897

I think both these test cases and be reduced. The select isn't needed (return/func are sufficiently picky about the type to be representative for this test) and the shape.const_shape that we don't fold is not needed after the select is removed.

906

random comment (unrelated to patch): I'm surprised that shape.const_shape [3] : tensor<1xindex> isn't rejected by the verifier. The ODS declaration has the return of Shape_ShapeOrExtentTensorType, which only allows !shape.shape or tensor<?xindex> (but seemingly not tensor<1xindex>).

[jurahul]

mlir/test/Dialect/Shape/canonicalize.mlir

906

That's because:

def Shape_ExtentTensorType :
    1DTensorOf<[Index]>,

Shape_ExtentTensorType allows any 1D tensor of index elements.

[silvas]

mlir/test/Dialect/Shape/canonicalize.mlir

906

Oh, I think I found it. Shape_ExtentTensorType inherits from 1DTensorOf<[Index]> but is BuildableType that creates tensor<?xindex>.

https://github.com/llvm/llvm-project/blob/df282215d497e15104ae9e182e083cdfa0bae3c2/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td#L110

That's kind of sketchy because BuildableType is supposed to be for types that only have one representation, which doesn't seem to be the case for Shape_ExtentTensorType if it also allows tensor<Nxindex> along with tensor<?xindex>.

https://mlir.llvm.org/docs/OpDefinitions/#type-inference

Buildable Types
Some type constraints may only have one representation, allowing for them to be directly buildable; for example the I32 or Index types. Types in ODS may mark themselves as buildable by setting the builderCall field or inheriting from the BuildableType class.

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

 def IndexToSizeToIndexCanonicalization : Pat<
   (Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
   (replaceWithValue $arg)>;
 
 def SizeToIndexToSizeCanonicalization : Pat<
   (Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
   (replaceWithValue $arg)>;
 
+// Fold tensor_cast(const_shape) to const_shape. This changes the type of
+// const_shape to the destination type of the cast.
 def TensorCastConstShape : Pat <
-  (TensorCastOp (Shape_ConstShapeOp:$c $ty)), (replaceWithValue $c)>;
+  (TensorCastOp (Shape_ConstShapeOp $arg)), (Shape_ConstShapeOp $arg)>;

mlir/test/Dialect/Shape/canonicalize.mlir

func @fold_assuming_all_single_element(%arg: tensor<?xindex>) {
   %0 = "test.source"() : () -> (!shape.witness)
   %1 = shape.assuming_all %0
   "consume.witness"(%1) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
-// Fold tensor_cast of a const_shape to const_shape
-// CHECK-LABEL: @fold_tensor_cast_of_const_shape
-func @fold_tensor_cast_of_const_shape(%arg: tensor<?xindex>) {
+// Verify that tensor_cast folding uses the correct type
+// CHECK-LABEL: @fold_tensor_cast_of_const_shape_returned
+func @fold_tensor_cast_of_const_shape_returned(%arg: i1) -> tensor<1xindex> {
+  // CHECK: constant dense<2> : tensor<1xindex>
   // CHECK-NOT: tensor_cast
   %0 = shape.const_shape [2] : tensor<?xindex>
   %1 = tensor_cast %0 : tensor<?xindex> to tensor<1xindex>
-  %2 = shape.cstr_broadcastable %1, %0 : tensor<1xindex>, tensor<?xindex>
-  "consume.witness"(%2) : (!shape.witness) -> ()
-  return
+  return %1 : tensor<1xindex>
+}
+
+// -----
+
+// Verify that tensor_cast folding uses the correct type
+// CHECK-LABEL: @fold_tensor_cast_of_const_shape_returned_dynamic
+func @fold_tensor_cast_of_const_shape_returned_dynamic(%arg: i1) -> tensor<?xindex> {
+  // CHECK: shape.const_shape [2] : tensor<?xindex>
+  // CHECK-NOT: tensor_cast
+  %0 = shape.const_shape [2] : tensor<1xindex>
+  %1 = tensor_cast %0 : tensor<1xindex> to tensor<?xindex>
+  return %1 : tensor<?xindex>
 }