[mlir] Don't allow dynamic extent tensor types for ConstShapeOp.

[akuegel]

Summary

ConstShapeOp has a constant shape, so its type can always be static.
We still allow it to have ShapeType though.

[akuegel]

Rebase.

[frgossen]

I agree, dynamically shaped constants seem counter-intuitive. They may be needed for folding ops w/o changing their result type though.

[akuegel]

Sean, Frederik told me that due to folding, we may want to still allow dynamic extent tensors, and that you initially suggested such a pattern. Therefore asking for your opinion on this change :)
With this change, I was able to find a lowering from Tensorflow dialect to HLO dialect where we did create dynamic extent tensors although it would have been possible to create static extent tensors. And also the bug fixed in https://reviews.llvm.org/D111127 would not have happened if we don't allow dynamic extent tensors for ConstantShapeOp.

[silvas]

I think it's fine for const_shape to always be static and require a tensor.cast to erase the static shape to dynamic.

[akuegel]

Thanks Sean, then I guess we are in agreement that this change might make sense. I guess we can give this change a try, and if it turns out that we actually still better want to allow dynamic type for const_shape, I can revert it?
Anyone wants to approve?

[frgossen]

Sg.

mlir/lib/Dialect/Shape/IR/Shape.cpp

 bool mlir::shape::ConstShapeOp::isCompatibleReturnTypes(TypeRange l,
                                                         TypeRange r) {
   if (l.size() != 1 || r.size() != 1)
     return false;
 
   Type lhs = l.front();
   Type rhs = r.front();
 
-  if (lhs == rhs)
-    return true;
-
   if (lhs.isa<ShapeType>() || rhs.isa<ShapeType>())
     // Shape type is compatible with all other valid return types.
     return true;
-
-  return succeeded(verifyCompatibleShapes(lhs, rhs));
+  return lhs == rhs;
 }
 
 //===----------------------------------------------------------------------===//
 // CstrBroadcastableOp
 //===----------------------------------------------------------------------===//
 
 void CstrBroadcastableOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {

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

 include "mlir/Dialect/Shape/IR/ShapeOps.td"
 include "mlir/Dialect/StandardOps/IR/Ops.td"
 include "mlir/Dialect/Tensor/IR/TensorOps.td"
 
 def AllInputShapesEq : Constraint<CPred< [{
   llvm::all_of($0, [&](mlir::Value val) {
     return $0[0] == val;
   })
 }]>>;
 
 def HasSingleElement : Constraint<CPred< [{
   $0.size() == 1
 }]>>;
 
+def HasStaticShape : Constraint<CPred< [{
+  $0.getType().dyn_cast<ShapedType>().hasStaticShape()
+}]>>;
+
 // Canonicalization patterns.
 
 def AssumingAllOneOp : Pat<(Shape_AssumingAllOp $args),
                            (replaceWithValue $args),
                            [(HasSingleElement $args)]>;
 
 def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),

 
 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 <
-  (Tensor_CastOp (Shape_ConstShapeOp $arg)), (Shape_ConstShapeOp $arg)>;
+  (Tensor_CastOp:$res (Shape_ConstShapeOp $arg)), (Shape_ConstShapeOp $arg),
+  [(HasStaticShape $res)]>;

mlir/test/Dialect/Shape/canonicalize.mlir

 // RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize %s | FileCheck %s
 
 // CHECK-LABEL: func @f
-func @f(%arg0: tensor<2x3x4xf32>) -> tensor<?xindex> {
-  // CHECK: shape.const_shape [2, 3, 4] : tensor<?xindex>
-  %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<?xindex>
-  return %0 : tensor<?xindex>
+func @f(%arg0: tensor<2x3x4xf32>) -> tensor<3xindex> {
+  // CHECK: shape.const_shape [2, 3, 4] : tensor<3xindex>
+  %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<3xindex>
+  return %0 : tensor<3xindex>
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK-DAG: shape.const_shape [2, 3] : !shape.shape
func @f() -> !shape.shape {
   %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
 
 // Basic case including extent tensors.
 // CHECK-LABEL: @broadcast
-func @broadcast() -> tensor<?xindex> {
-  // CHECK: shape.const_shape [7, 2] : tensor<?xindex>
-  %0 = shape.const_shape [1, 2] : tensor<?xindex>
-  %1 = shape.const_shape [7, 1] : tensor<?xindex>
+func @broadcast() -> tensor<2xindex> {
+  // CHECK: shape.const_shape [7, 2] : tensor<2xindex>
+  %0 = shape.const_shape [1, 2] : tensor<2xindex>
+  %1 = shape.const_shape [7, 1] : tensor<2xindex>
   %2 = shape.broadcast %0, %1
-      : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
-  return %2 : tensor<?xindex>
+      : tensor<2xindex>, tensor<2xindex> -> tensor<2xindex>
+  return %2 : tensor<2xindex>
 }
 
 // -----
 
 // Basic case including extent tensors.
 // CHECK-LABEL: @broadcast
 func @broadcast() -> !shape.shape {
   // CHECK: shape.const_shape [7, 2] : !shape.shape
-  %0 = shape.const_shape [1, 2] : tensor<?xindex>
-  %1 = shape.const_shape [7, 1] : tensor<?xindex>
-  %2 = shape.broadcast %0, %1 : tensor<?xindex>, tensor<?xindex> -> !shape.shape
+  %0 = shape.const_shape [1, 2] : tensor<2xindex>
+  %1 = shape.const_shape [7, 1] : tensor<2xindex>
+  %2 = shape.broadcast %0, %1 : tensor<2xindex>, tensor<2xindex> -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
 
 // Rhs is a scalar.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {

 }
 
 // -----
 
 // Basic folding.
 // CHECK-LABEL: func @basic
 func @basic() -> index {
   // CHECK: constant 2 : index
-  %0 = shape.const_shape [0, 1, 2] : tensor<?xindex>
+  %0 = shape.const_shape [0, 1, 2] : tensor<3xindex>
   %c2 = constant 2 : index
-  %1 = shape.get_extent %0, %c2 : tensor<?xindex>, index -> index
+  %1 = shape.get_extent %0, %c2 : tensor<3xindex>, index -> index
   return %1 : index
 }
 
 // -----
 
 // Should not fold.
 // CHECK-LABEL: func @out_of_bounds
 func @out_of_bounds() -> index {
   // CHECK: shape.const_shape
   // CHECK: shape.get_extent
-  %0 = shape.const_shape [0, 1, 2] : tensor<?xindex>
+  %0 = shape.const_shape [0, 1, 2] : tensor<3xindex>
   %c3 = constant 3 : index
-  %1 = shape.get_extent %0, %c3 : tensor<?xindex>, index -> index
+  %1 = shape.get_extent %0, %c3 : tensor<3xindex>, index -> index
   return %1 : index
 }
 
 // -----
 
 // Should not fold.
 // CHECK-LABEL: func @not_const
 func @not_const(%arg0: tensor<?xindex>) -> index {
func @f(%arg : !shape.shape) -> !shape.shape {
   %1 = shape.any %0, %arg : !shape.shape, !shape.shape -> !shape.shape
   return %1 : !shape.shape
 }
 
 // -----
 
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
-func @f(%arg : tensor<?xindex>) -> tensor<?xindex> {
-  // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<?xindex>
-  // CHECK-NEXT: return %[[CS]] : tensor<?xindex>
-  %0 = shape.const_shape [2, 3, 4] : tensor<?xindex>
-  %1 = shape.any %0, %arg : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
-  return %1 : tensor<?xindex>
+func @f(%arg : tensor<?xindex>) -> tensor<3xindex> {
+  // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<3xindex>
+  // CHECK-NEXT: return %[[CS]] : tensor<3xindex>
+  %0 = shape.const_shape [2, 3, 4] : tensor<3xindex>
+  %1 = shape.any %0, %arg : tensor<3xindex>, tensor<?xindex> -> tensor<3xindex>
+  return %1 : tensor<3xindex>
 }
 
 // -----
 
 // Folding of any with partially constant operands is not yet implemented.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.any

 
 // -----
 
 // Fold `rank` based on constant extent tensor.
 // CHECK-LABEL: @fold_rank
 func @fold_rank() -> index {
   // CHECK: %[[RESULT:.*]] = constant 5 : index
   // CHECK: return %[[RESULT]] : index
-  %shape = shape.const_shape [3, 4, 5, 6, 7] : tensor<?xindex>
-  %rank = shape.rank %shape : tensor<?xindex> -> index
+  %shape = shape.const_shape [3, 4, 5, 6, 7] : tensor<5xindex>
+  %rank = shape.rank %shape : tensor<5xindex> -> index
   return %rank : index
 }
 
 // -----
 
 // Do not fold `rank` for non-constant extent tensors.
 // CHECK-LABEL: @dont_fold_rank
 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index

 // -----
 
 // Fold `shape_eq` for equal and constant shapes.
 // CHECK-LABEL: @shape_eq_fold_1
 func @shape_eq_fold_1() -> i1 {
   // CHECK: %[[RESULT:.*]] = constant true
   // CHECK: return %[[RESULT]] : i1
   %a = shape.const_shape [1, 2, 3] : !shape.shape
-  %b = shape.const_shape [1, 2, 3] : tensor<?xindex>
-  %c = shape.const_shape [1, 2, 3] : tensor<?xindex>
-  %result = shape.shape_eq %a, %b, %c : !shape.shape, tensor<?xindex>, tensor<?xindex>
+  %b = shape.const_shape [1, 2, 3] : tensor<3xindex>
+  %c = shape.const_shape [1, 2, 3] : tensor<3xindex>
+  %result = shape.shape_eq %a, %b, %c : !shape.shape, tensor<3xindex>, tensor<3xindex>
   return %result : i1
 }
 
 // -----
 
 // Fold `shape_eq` for different but constant shapes of same length.
 // CHECK-LABEL: @shape_eq_fold_0
 func @shape_eq_fold_0() -> i1 {
   // CHECK: %[[RESULT:.*]] = constant false
   // CHECK: return %[[RESULT]] : i1
-  %a = shape.const_shape [1, 2, 3] : tensor<?xindex>
-  %b = shape.const_shape [4, 5, 6] : tensor<?xindex>
-  %c = shape.const_shape [4, 5, 6] : tensor<?xindex>
-  %result = shape.shape_eq %a, %b, %c : tensor<?xindex>, tensor<?xindex>, tensor<?xindex>
+  %a = shape.const_shape [1, 2, 3] : tensor<3xindex>
+  %b = shape.const_shape [4, 5, 6] : tensor<3xindex>
+  %c = shape.const_shape [4, 5, 6] : tensor<3xindex>
+  %result = shape.shape_eq %a, %b, %c : tensor<3xindex>, tensor<3xindex>, tensor<3xindex>
   return %result : i1
 }
 
 // -----
 
 // Fold `shape_eq` for different but constant shapes of different length.
 // CHECK-LABEL: @shape_eq_fold_0
 func @shape_eq_fold_0() -> i1 {

 
 // -----
 
 // 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: shape.const_shape [2] : tensor<1xindex>
   // CHECK-NOT: tensor.cast
-  %0 = shape.const_shape [2] : tensor<?xindex>
-  %1 = tensor.cast %0 : tensor<?xindex> to tensor<1xindex>
+  %0 = shape.const_shape [2] : tensor<1xindex>
+  %1 = tensor.cast %0 : tensor<1xindex> to tensor<1xindex>
   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
+// CHECK-LABEL: @dont_fold_tensor.cast_of_const_shape_returned_dynamic
+func @dont_fold_tensor.cast_of_const_shape_returned_dynamic(%arg: i1) -> tensor<?xindex> {
+  // CHECK: %[[CONST_SHAPE:.*]] = shape.const_shape [2] : tensor<1xindex>
+  // CHECK: tensor.cast %[[CONST_SHAPE]] : tensor<1xindex> to tensor<?xindex>
   %0 = shape.const_shape [2] : tensor<1xindex>
   %1 = tensor.cast %0 : tensor<1xindex> to tensor<?xindex>
   return %1 : tensor<?xindex>
 }
 
 // -----
 
 // CHECK-LABEL: @is_broadcastable_on_same_shape

mlir/test/Dialect/Shape/ops.mlir

func @test_shape_num_elements_unknown() {
   %0 = "shape.unknown_shape"() : () -> !shape.shape
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %2 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @const_shape() {
   %0 = shape.const_shape [1, 2, 3] : !shape.shape
-  %1 = shape.const_shape [4, 5, 6] : tensor<?xindex>
   %2 = shape.const_shape [4, 5, 6] : tensor<3xindex>
   return
 }
 
 func @test_shape_num_elements_fixed() {
   %0 = shape.const_shape [1, 57, 92] : !shape.shape
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @test_broadcast_fixed() {
   %0 = shape.const_shape [10, 1, 57, 92] : !shape.shape
   %1 = shape.const_shape [4, 57, 92] : !shape.shape
   %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
-func @test_broadcast_extents() -> tensor<?xindex> {
-  %0 = shape.const_shape [10, 1, 57, 92] : tensor<?xindex>
-  %1 = shape.const_shape [4, 57, 92] : tensor<?xindex>
-  %2 = shape.broadcast %0, %1 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
-  return %2 : tensor<?xindex>
+func @test_broadcast_extents() -> tensor<4xindex> {
+  %0 = shape.const_shape [10, 1, 57, 92] : tensor<4xindex>
+  %1 = shape.const_shape [4, 57, 92] : tensor<3xindex>
+  %2 = shape.broadcast %0, %1 : tensor<4xindex>, tensor<3xindex> -> tensor<4xindex>
+  return %2 : tensor<4xindex>
 }
 
 func @test_shape_any_fixed() {
   %0 = shape.const_shape [4, 57, 92] : !shape.shape
   %1 = shape.const_shape [4, 57, 92] : !shape.shape
   %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
func @test_shape_any_fixed_mismatch() {
   %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_parse_const_shape() {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
-  %2 = shape.const_shape [1, 2, 3] : tensor<?xindex>
+  %2 = shape.const_shape [1, 2, 3] : tensor<3xindex>
   return
 }
 
 func @test_shape_of(%arg0: tensor<?xf32>) -> tensor<?xindex> {
   %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
 
func @get_extent_on_mixed_operands(%arg : tensor<?xindex>) -> !shape.size {
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> !shape.size
   return %result : !shape.size
 }
 
 func @any() {
   %0 = shape.const_shape [1, 2, 3] : !shape.shape
   %1 = shape.const_shape [4, 5, 6] : !shape.shape
   %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
-  %3 = shape.const_shape [1, 2, 3] : tensor<?xindex>
-  %4 = shape.const_shape [4, 5, 6] : tensor<?xindex>
-  %5 = "shape.any"(%3, %4) : (tensor<?xindex>, tensor<?xindex>) -> tensor<?xindex>
+  %3 = shape.const_shape [1, 2, 3] : tensor<3xindex>
+  %4 = shape.const_shape [4, 5, 6] : tensor<3xindex>
+  %5 = "shape.any"(%3, %4) : (tensor<3xindex>, tensor<3xindex>) -> tensor<3xindex>
   return
 }
 
 func @num_elements_extent_tensor(%arg : tensor<?xindex>) -> index {
   %result = shape.num_elements %arg : tensor<?xindex> -> index
   return %result : index
 }