[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

36

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

36

$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

35

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

mlir/test/Dialect/Shape/canonicalize.mlir

886

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.

895

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

895

That's because:

def Shape_ExtentTensorType :
    1DTensorOf<[Index]>,

Shape_ExtentTensorType allows any 1D tensor of index elements.

[silvas]

mlir/test/Dialect/Shape/canonicalize.mlir

895

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)>;
 
 def TensorCastConstShape : Pat <
-  (TensorCastOp (Shape_ConstShapeOp:$c $ty)), (replaceWithValue $c)>;
+  (TensorCastOp (Shape_ConstShapeOp $ty)), (Shape_ConstShapeOp $ty)>;

mlir/test/Dialect/Shape/canonicalize.mlir

 func @fold_tensor_cast_of_const_shape(%arg: tensor<?xindex>) {
   // 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
 }
+
+// -----
+
+// Verify that tensor_cast folding uses the correct type
+// CHECK-LABEL: @fold_tensor_cast_of_const_shape_to_select
+func @fold_tensor_cast_of_const_shape_to_select(%arg: i1) -> tensor<1xindex> {
+  // CHECK-DAG: constant dense<2> : tensor<1xindex>
+  // CHECK-DAG: constant dense<3> : tensor<1xindex>
+  // CHECK-NOT: tensor_cast
+  %0 = shape.const_shape [2] : tensor<?xindex>
+  %1 = shape.const_shape [3] : tensor<1xindex>
+  %2 = tensor_cast %0 : tensor<?xindex> to tensor<1xindex>
+  %3 = select %arg, %1, %2 : i1, tensor<1xindex>
+  return %3 : tensor<1xindex>
+}
+
+// -----
+
+// Verify that tensor_cast folding uses the correct type
+// CHECK-LABEL: @fold_tensor_cast_of_const_shape_to_select_dynamic
+func @fold_tensor_cast_of_const_shape_to_select_dynamic(%arg: i1) -> tensor<?xindex> {
+  // CHECK-DAG: shape.const_shape [2] : tensor<?xindex>
+  // CHECK-DAG: shape.const_shape [3] : tensor<?xindex>
+  // CHECK-NOT: tensor_cast
+  %0 = shape.const_shape [2] : tensor<1xindex>
+  %1 = shape.const_shape [3] : tensor<?xindex>
+  %2 = tensor_cast %0 : tensor<1xindex> to tensor<?xindex>
+  %3 = select %arg, %1, %2 : i1, tensor<?xindex>
+  return %3 : tensor<?xindex>
+}