[MLIR][Shape] Generalize `shape.const_shape` to extent tensors

[frgossen]

Summary

The operation shape.const_shape was used for constants of type shape only.
We can now also use it to create constant extent tensors.

Depends On D84156

[herhut]

Shouldn't shape.const_shape always return a tensor? It cannot return an error currently, right?

[frgossen]

Not at the moment but I think eventually it should also be possible to create constant error shapes.
We can split const_shape into const_shape and const_error later if we want that.
If the return type of const_shape was constrained to an extent tensor at this point then we'd have to update all the operations before (because it's all over in the tests).

[frgossen]

Add test case

[silvas]

What do you think about

%1 = shape.const_shape [1, 2, 3]
%2 = shape.const_extent_tensor [4, 5, 6]

instead of

%1 = shape.const_shape [1, 2, 3] : !shape.shape
%2 = shape.const_shape [4, 5, 6] : tensor<?xindex>

?

That seems more succinct.

Also, the right thing here, given our current rationale, is to always return an extent tensor from this op since it cannot be an error. Are you planning on doing that migration as well? It seems that this patch puts us in an awkward intermediate state.

[stellaraccident]

mlir/test/Dialect/Shape/ops.mlir

84

What is the rationale for having a shape-specific constant op for this? Since tensor<?xindex> is a standard type, [std.]constant is the right way to be constructing constants for it, no?

[frgossen]

This is indeed an intermediate state.

We would like to update all operations so that they accept extent tensors. At that point, the operation const_shape could be split in two operations: const_shape for extent tensors and const_error.

Eventually, we would like to replace const_shape with std.constant. However, that requires type refinement as we can only create statically shaped extent tensors this way. When fed into a return operation, e.g., that causes problems otherwise.

mlir/test/Dialect/Shape/ops.mlir

84

That is what we wanted but it turns out that std.constant cannot create dynamically shaped tensor constants like tensor<?xindex>. The example here, shape.const_shape [1, 2, 3], would always be of type tensor<3xindex> which is problematic when consumed by, e.g., a return operation. Also casting is not really an option for constant folding.

In the long run, we would still like to use std.constant for this. However, this required type refinement. Maybe we can discuss this in our meeting today.

[frgossen]

Rebase

[frgossen]

Make helper static

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_BroadcastOp : Shape_Op<"broadcast", [Commutative]> {
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
 
   let hasFolder = 1;
 }
 
 def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
-  let summary = "Creates a constant of !shape.shape type";
+  let summary = "Creates a constant shape or extent tensor";
   let description = [{
-    Creates a !shape.shape with rank given by the length of `shape` and with
-    dimension sizes given by the values of `shape`.
+    Creates a constant shape or extent tensor. The individual extents are given
+    as the `shape` attribute. The number of these values equals the shape's
+    rank.
 
     ```mlir
-    %0 = shape.const_shape []
-    %1 = shape.const_shape [1, 2, 3]
+    %0 = shape.const_shape [] : !shape.shape
+    %1 = shape.const_shape [1, 2, 3] : !shape.shape
+    %2 = shape.const_shape [4, 5, 6] : tensor<?xindex>
     ```
   }];
   let arguments = (ins IndexElementsAttr:$shape);
-  let results = (outs Shape_ShapeType:$result);
+  let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
   // TODO: Move this to main so that all shape ops implement these.
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let hasFolder = 1;
 }
 
 def Shape_ConstSizeOp : Shape_Op<"const_size", [

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

 
 using namespace mlir;
 using namespace mlir::shape;
 
 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
+static RankedTensorType getExtentTensorType(OpBuilder &builder) {
+  return RankedTensorType::get({ShapedType::kDynamicSize},
+                               builder.getIndexType());
+}
+
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
            WitnessType>();
   // Allow unknown operations during prototyping and testing. As the dialect is
   // still evolving it makes it simple to start with an unregistered ops and
   // try different variants before actually defining the op.
   allowUnknownOperations();
 }
 
 Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
                                              Attribute value, Type type,
                                              Location loc) {
-  if (auto shapeType = type.dyn_cast<ShapeType>())
+  if (type.isa<ShapeType>() || type == getExtentTensorType(builder))
     return builder.create<ConstShapeOp>(loc, type,
                                         value.cast<DenseIntElementsAttr>());
-  if (auto sizeType = type.dyn_cast<SizeType>())
+  if (type.isa<SizeType>())
     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
-  if (auto witnessType = type.dyn_cast<WitnessType>())
+  if (type.isa<WitnessType>())
     return builder.create<ConstWitnessOp>(loc, type, value.cast<BoolAttr>());
   return nullptr;
 }
 
 /// Parse a type registered to this dialect.
 Type ShapeDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))

 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
   p << "shape.const_shape ";
   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"shape"});
   p << "[";
   interleaveComma(op.shape().getValues<int64_t>(), p,
                   [&](int64_t i) { p << i; });
-  p << "]";
+  p << "] : ";
+  p.printType(op.getType());
 }
 
 static ParseResult parseConstShapeOp(OpAsmParser &parser,
                                      OperationState &result) {
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
   // We piggy-back on ArrayAttr parsing, though we don't internally store the
   // shape as an ArrayAttr.
static ParseResult parseConstShapeOp(OpAsmParser &parser,
   for (Attribute extent : extentsArray) {
     IntegerAttr attr = extent.dyn_cast<IntegerAttr>();
     if (!attr)
       return failure();
     ints.push_back(attr.getInt());
   }
   Builder &builder = parser.getBuilder();
   result.addAttribute("shape", builder.getIndexTensorAttr(ints));
-
-  result.types.push_back(ShapeType::get(builder.getContext()));
+  Type resultTy;
+  if (parser.parseColonType(resultTy))
+    return failure();
+  result.types.push_back(resultTy);
   return success();
 }
 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
 
 //===----------------------------------------------------------------------===//
 // CstrBroadcastableOp
 //===----------------------------------------------------------------------===//

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>) -> !shape.shape {
-  // CHECK: shape.const_shape [2, 3, 4]
+  // CHECK: shape.const_shape [2, 3, 4] : !shape.shape
   %0 = "shape.shape_of"(%arg0) : (tensor<2x3x4xf32>) -> !shape.shape
   return %0 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
-  // CHECK: shape.const_shape [2, 3]
-  // CHECK: shape.const_shape [4, 5]
+  // CHECK: shape.const_shape [2, 3] : !shape.shape
+  // CHECK: shape.const_shape [4, 5] : !shape.shape
   %c2 = constant 2 : i32
-  %0 = shape.const_shape [2, 3, 4, 5]
+  %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
   %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
   return %head, %tail : !shape.shape, !shape.shape
 
 }
 
 // -----
 
 // Negative split point.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
-  // CHECK: shape.const_shape [2, 3, 4]
-  // CHECK: shape.const_shape [5]
+  // CHECK: shape.const_shape [2, 3, 4] : !shape.shape
+  // CHECK: shape.const_shape [5] : !shape.shape
   %c-1 = constant -1 : i32
-  %0 = shape.const_shape [2, 3, 4, 5]
+  %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
   %head, %tail = "shape.split_at"(%0, %c-1) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
   return %head, %tail : !shape.shape, !shape.shape
 }
 
 // -----
 
 // Out of range split point. No folding.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK: shape.split_at
   %c5 = constant 5 : i32
-  %0 = shape.const_shape [2, 3, 4, 5]
+  %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
   %head, %tail = "shape.split_at"(%0, %c5) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
   return %head, %tail : !shape.shape, !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
-  // CHECK: shape.const_shape [7, 2]
-  %0 = shape.const_shape [1, 2]
-  %1 = shape.const_shape [7, 1]
+  // CHECK: shape.const_shape [7, 2] : !shape.shape
+  %0 = shape.const_shape [1, 2] : !shape.shape
+  %1 = shape.const_shape [7, 1] : !shape.shape
   %2 = shape.broadcast %0, %1
   return %2 : !shape.shape
 }
 
 // -----
 
 // Rhs is a scalar.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {
   // CHECK: return %arg0
-  %0 = shape.const_shape []
+  %0 = shape.const_shape [] : !shape.shape
   %1 = shape.broadcast %arg0, %0
   return %1 : !shape.shape
 }
 
 // -----
 
 // Lhs is a scalar.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {
   // CHECK: return %arg0
-  %0 = shape.const_shape []
+  %0 = shape.const_shape [] : !shape.shape
   %1 = shape.broadcast %0, %arg0
   return %1 : !shape.shape
 }
 
 // -----
 
 // Lhs is a scalar and rhs is constant.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
-  // CHECK: %[[CST:.*]] = shape.const_shape [1, 2, 3]
+  // CHECK: %[[CST:.*]] = shape.const_shape [1, 2, 3] : !shape.shape
   // CHECK: return %[[CST]]
-  %0 = shape.const_shape []
-  %1 = shape.const_shape [1, 2, 3]
+  %0 = shape.const_shape [] : !shape.shape
+  %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %2 = shape.broadcast %0, %1
   return %2 : !shape.shape
 }
 
 // -----
 
 // Incompatible shapes. No folding.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.broadcast
-  %0 = shape.const_shape [2]
-  %1 = shape.const_shape [7]
+  %0 = shape.const_shape [2] : !shape.shape
+  %1 = shape.const_shape [7] : !shape.shape
   %2 = shape.broadcast %0, %1
   return %2 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
-  // CHECK: shape.const_shape [0, 1, 2, 3]
-  %lhs = shape.const_shape [0, 1]
-  %rhs = shape.const_shape [2, 3]
+  // CHECK: shape.const_shape [0, 1, 2, 3] : !shape.shape
+  %lhs = shape.const_shape [0, 1] : !shape.shape
+  %rhs = shape.const_shape [2, 3] : !shape.shape
   %0 = shape.concat %lhs, %rhs
   return %0 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> tensor<2xindex> {
   // CHECK: constant dense<[0, 1]> : tensor<2xindex>
-  %cs = shape.const_shape [0, 1]
+  %cs = shape.const_shape [0, 1] : !shape.shape
   %0 = shape.to_extent_tensor %cs : tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f()
 func @f() -> !shape.shape {
-  // CHECK: shape.const_shape [3, 5, 11]
+  // CHECK: shape.const_shape [3, 5, 11] : !shape.shape
   %e0 = constant 3 : index
   %e1 = constant 5 : index
   %e2 = constant 11 : index
   %ret = shape.from_extents %e0, %e1, %e2
   return %ret : !shape.shape
 }
 
 // -----

 }
 
 // -----
 
 // Fold number of elements computation.
 // CHECK-LABEL: func @num_elements
 func @num_elements() -> !shape.size {
   // CHECK-NOT: shape.const_shape
-  %shape = shape.const_shape [4, 5, 6]
+  %shape = shape.const_shape [4, 5, 6] : !shape.shape
   // CHECK-NOT: shape.num_elements
   %num_elements = shape.num_elements %shape
   // CHECK: %[[NUM:.*]] = shape.const_size 120
   // CHECK-NEXT: return %[[NUM]] : !shape.size
   return %num_elements : !shape.size
 }
 
 // -----
 
 // No folding.
 // CHECK-LABEL: func @nonfoldable_num_elements
 func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {
   // CHECK-NOT: shape.const_{{.*}}
   %num_elements = shape.num_elements %shape
   return %num_elements : !shape.size
 }
 
 // -----
 
 // Basic folding.
 // CHECK-LABEL: func @basic
 func @basic() -> !shape.size {
   // CHECK: shape.const_size 2
-  %0 = shape.const_shape [0, 1, 2]
+  %0 = shape.const_shape [0, 1, 2] : !shape.shape
   %c2 = shape.const_size 2
   %1 = shape.get_extent %0, %c2
   return %1 : !shape.size
 }
 
 // -----
 
 // Should not fold.
 // CHECK-LABEL: func @out_of_bounds
 func @out_of_bounds() -> !shape.size {
   // CHECK: shape.const_shape
   // CHECK: shape.get_extent
-  %0 = shape.const_shape [0, 1, 2]
+  %0 = shape.const_shape [0, 1, 2] : !shape.shape
   %c3 = shape.const_size 3
   %1 = shape.get_extent %0, %c3
   return %1 : !shape.size
 }
 
 // -----
 
 // Should not fold.

 
 // -----
 // cstr_eq with equal const_shapes can be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [0, 1]
-  %cs1 = shape.const_shape [0, 1]
-  %cs2 = shape.const_shape [0, 1]
+  %cs0 = shape.const_shape [0, 1] : !shape.shape
+  %cs1 = shape.const_shape [0, 1] : !shape.shape
+  %cs2 = shape.const_shape [0, 1] : !shape.shape
   %0 = shape.cstr_eq %cs0, %cs1, %cs2
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // cstr_eq with unequal const_shapes cannot be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.cstr_eq
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [0, 1]
-  %cs1 = shape.const_shape [3, 1]
+  %cs0 = shape.const_shape [0, 1] : !shape.shape
+  %cs1 = shape.const_shape [3, 1] : !shape.shape
   %0 = shape.cstr_eq %cs0, %cs1
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // cstr_eq without const_shapes cannot be folded
 // CHECK-LABEL: func @f

 
 // -----
 
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
   // CHECK-NEXT: return %[[CS]]
-  %0 = shape.const_shape [2, 3, 4]
+  %0 = shape.const_shape [2, 3, 4] : !shape.shape
   %1 = shape.any %0, %arg0
   return %1 : !shape.shape
 }
 
 
 // -----
 
 // Folding of any with partially constant operands is not yet implemented.

 
 // -----
 // Broadcastable with broadcastable constant shapes can be removed.
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [3, 1]
-  %cs1 = shape.const_shape [1, 5]
+  %cs0 = shape.const_shape [3, 1] : !shape.shape
+  %cs1 = shape.const_shape [1, 5] : !shape.shape
   %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with non-broadcastable constant shapes is always false
 // CHECK-LABEL: func @static_non_broadcastable
 func @static_non_broadcastable() {
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [1, 3]
-  %cs1 = shape.const_shape [1, 5]
+  %cs0 = shape.const_shape [1, 3] : !shape.shape
+  %cs1 = shape.const_shape [1, 5] : !shape.shape
   %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable without guaranteed broadcastable shapes cannot be removed.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_shape
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %cs0 = shape.const_shape [1,3]
+  %cs0 = shape.const_shape [1, 3] : !shape.shape
   %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
 // Broadcastable with non-constant but known equal shapes can be removed.
 // CHECK-LABEL: func @f

 
 // -----
 
 // Fold `rank` based on constant shape.
 // CHECK-LABEL: @fold_rank
 func @fold_rank() -> !shape.size {
   // CHECK-DAG: %[[RESULT:.*]] = shape.const_size 5
   // CHECK-DAG: return %[[RESULT]] : !shape.size
-  %shape = shape.const_shape [3, 4, 5, 6, 7]
+  %shape = shape.const_shape [3, 4, 5, 6, 7] : !shape.shape
   %rank = shape.rank %shape : !shape.shape
   return %rank : !shape.size
 }
 
 // -----
 
 // Do not fold `rank` if shape is dynamic.
 // CHECK-LABEL: @dont_fold_rank

 // -----
 
 // Canonicalize scalar cstr_broadcastable checks
 // CHECK-LABEL: @cstr_broadcastable_scalar
 func @cstr_broadcastable_scalar(%arg0 : tensor<?xf32>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.const_shape []
+  %0 = shape.const_shape [] : !shape.shape
   %1 = shape.shape_of %arg0 : tensor<?xf32>
   %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 

 
 // -----
 
 // 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]
-  %b = shape.const_shape [1, 2, 3]
-  %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
+  %a = shape.const_shape [1, 2, 3] : !shape.shape
+  %b = shape.const_shape [1, 2, 3] : tensor<?xindex>
+  %result = shape.shape_eq %a, %b : !shape.shape, tensor<?xindex>
   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]
-  %b = shape.const_shape [4, 5, 6]
-  %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
+  %a = shape.const_shape [1, 2, 3] : tensor<?xindex>
+  %b = shape.const_shape [4, 5, 6] : tensor<?xindex>
+  %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
   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 {
   // CHECK: %[[RESULT:.*]] = constant false
   // CHECK: return %[[RESULT]] : i1
-  %a = shape.const_shape [1, 2, 3, 4, 5, 6]
-  %b = shape.const_shape [1, 2, 3]
+  %a = shape.const_shape [1, 2, 3, 4, 5, 6] : !shape.shape
+  %b = shape.const_shape [1, 2, 3] : !shape.shape
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }
 
 // -----
 
 // Do not fold `shape_eq` for non-constant shapes.
 // CHECK-LABEL: @shape_eq_do_not_fold
 // CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i1
 func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 {
   // CHECK: %[[B:.*]] = shape.const_shape [4, 5, 6]
   // CHECK: %[[RESULT:.*]] = shape.shape_eq %[[A]], %[[B]] : !shape.shape, !shape.shape
   // CHECK: return %[[RESULT]] : i1
-  %b = shape.const_shape [4, 5, 6]
+  %b = shape.const_shape [4, 5, 6] : !shape.shape
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }

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>
+  return
+}
+
 func @test_shape_num_elements_fixed() {
-  %0 = shape.const_shape [1, 57, 92]
+  %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]
-  %1 = shape.const_shape [4, 57, 92]
+  %0 = shape.const_shape [10, 1, 57, 92] : !shape.shape
+  %1 = shape.const_shape [4, 57, 92] : !shape.shape
   %2 = shape.broadcast %0, %1
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_fixed() {
-  %0 = shape.const_shape [4, 57, 92]
-  %1 = shape.const_shape [4, 57, 92]
+  %0 = shape.const_shape [4, 57, 92] : !shape.shape
+  %1 = shape.const_shape [4, 57, 92] : !shape.shape
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_unknown() {
-  %0 = shape.const_shape [4, -1, 92]
-  %1 = shape.const_shape [-1, 57, 92]
+  %0 = shape.const_shape [4, -1, 92] : !shape.shape
+  %1 = shape.const_shape [-1, 57, 92] : !shape.shape
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_fixed_mismatch() {
-  %0 = shape.const_shape [4, 57, 92]
-  %1 = shape.const_shape [2, 57, 92]
+  %0 = shape.const_shape [4, 57, 92] : !shape.shape
+  %1 = shape.const_shape [2, 57, 92] : !shape.shape
   %2 = "shape.join"(%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 []
-  %1 = shape.const_shape [1, 2, 3]
+  %0 = shape.const_shape [] : !shape.shape
+  %1 = shape.const_shape [1, 2, 3] : !shape.shape
+  %2 = shape.const_shape [1, 2, 3] : tensor<?xindex>
   return
 }
 
 func @test_shape_of(%arg0: tensor<?xf32>) -> !shape.shape {
   %0 = shape.shape_of %arg0 : tensor<?xf32>
   return %0 : !shape.shape
 }
 
 func @test_constraints() {
-  %0 = shape.const_shape []
-  %1 = shape.const_shape [1, 2, 3]
+  %0 = shape.const_shape [] : !shape.shape
+  %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   %w1 = shape.cstr_eq %0, %1
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false
   %w4 = shape.assuming_all %w0, %w1, %w2, %w3
   shape.assuming %w4 -> !shape.shape {
     %2 = shape.any %0, %1
     shape.assuming_yield %2 : !shape.shape