[MLIR][Shape] Allow for `shape_of` to return extent tensors

[frgossen]

Summary

The operation shape.shape_of now returns an extent tensor tensor<?xindex> in cases when no error are possible. All consuming operation will eventually accept both, shapes and extent tensors.

Depends On D84158

[pifon2a]

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

25–26

nit: move inline before RankedTensorType

[herhut]

We should discuss whether we in the long run want to allow known-size tensors for known-rank inputs.

[jpienaar]

In D84160#2161985, @herhut wrote:

We should discuss whether we in the long run want to allow known-size tensors for known-rank inputs.

We should yes (that does make the above more complicated)

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

395

I don't understand the tuple part, did you mean type?

396

Doesn't it just return the shape?

397

But this also can take a ValueShapeType as input and that can be an error already (e.g., you can invoke one shape func from another, we need to discuss the create function there we haven't gotten to it in the last couple of syncs, I should try and implement today so that we can discuss with code).

Only when input is just a ShapedType is no error possible. It would therefore seem that this one has:

ValueShapeType -> !shape.shape
ShapedType -> tensor<?xindex>

With it being legal but canonicalized away to have ShapedType -> !shape.shape during lowering.

[rriddle]

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

25–26

Why is this being marked as inline? Also, static functions should not be in anonymous namespaces and should be marked as 'static' instead.

[frgossen]

Address comments

[frgossen]

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

395

The documentation for value_shape says "Conceptually this is a tuple of a value (potentially unknown) and shape.type". I thought calling it a value-shape tuple would help in the description. Apparently it doesn't.

396

It can return a shape (shape.shape) or an extent tensor (tensor<?xindex>)

397

I did not fully understand the value_shape it seems. The operation now allows for both variants.
Would the variant "ShapedType -> !shape.shape" ever be useful? Eventually, all operations will accept extent tensors.

[jpienaar]

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

397

Only as intermediate form really that I can think of. And even there it is dubious, as it would reified with the correct type ...

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

655

Up to you, but I'd relax this one as the above loses info, while this one is just overly broad. BUT I can't really think in practice where this would happen - I can think of convoluted ways it could happen ... But none that I see us using.

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

let typeDescription = [{
     `shape.value_shape` represents the value produced by an operation (this
     corresponds to `Value` in the compiler) and a shape. Conceptually this is a
     tuple of a value (potentially unknown) and `shape.type`. The value and shape
     can either or both be unknown. If both the `value` and `shape` are known,
     then the shape of `value` is conformant with `shape`.
   }];
 }
 
+def Shape_ExtentTensorType :
+    1DTensorOf<[Index]>,
+    BuildableType<"::mlir::RankedTensorType::get({ShapedType::kDynamicSize}, "
+                  "$_builder.getType<::mlir::IndexType>())"> {
+  let typeDescription = [{
+    The extent tensor is a tensor of rank one with arbitrarily many index
+    elements. Like `!shape.shape`, it is used to represent shapes with the
+    difference that it is guaranteed to be error-free.
+  }];
+}
+
 def Shape_ShapeOrSizeType : AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
   "shape or size">;
 
-def Shape_ExtentTensorType : 1DTensorOf<[Index]>;
-
 def Shape_ShapeOrExtentTensorType : AnyTypeOf<[Shape_ShapeType,
                                                Shape_ExtentTensorType],
                                               "shape or extent tensor">;
 
 #endif // SHAPE_BASE_TD

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

def Shape_ReduceOp : Shape_Op<"reduce",
   let verifier = [{ return ::verify(*this); }];
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 def Shape_ShapeOfOp : Shape_Op<"shape_of", [NoSideEffect]> {
   let summary = "Returns shape of a value or shaped type operand";
 
+  let description = [{
+    The operation takes a value or a shaped operand as an argument and it
+    returns a shape or extent tensor.
+  }];
+
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$arg);
-  let results = (outs Shape_ShapeType:$result);
+  let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
-  let assemblyFormat = "$arg `:` type($arg) attr-dict";
+  let assemblyFormat = "$arg `:` type($arg) `->` type($result) attr-dict";
 
+  let verifier = [{ return ::verify(*this); }];
   let hasFolder = 1;
 }
 
 def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let summary = "Casts between index types of the shape and standard dialect";
   let description = [{
     Converts a `shape.size` to a standard index.
     This operation and its inverse, `index_to_size`, facilitate index conversion

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

for (int64_t i = 0; i < rank; i++) {
         dimValues.push_back(dimVal);
       } else {
         int64_t dim = rankedTensorTy.getDimSize(i);
         auto dimVal = rewriter.create<ConstantIndexOp>(loc, dim);
         dimValues.push_back(dimVal);
       }
     }
 
-    // Materialize shape as ranked tensor.
-    rewriter.replaceOpWithNewOp<TensorFromElementsOp>(op.getOperation(),
-                                                      dimValues);
+    // Materialize extent tensor.
+    Value staticExtentTensor =
+        rewriter.create<TensorFromElementsOp>(loc, dimValues);
+    rewriter.replaceOpWithNewOp<TensorCastOp>(op, staticExtentTensor,
+                                              op.getType());
     return success();
   }
 };
 
 class ConstSizeOpConverter : public OpConversionPattern<ConstSizeOp> {
 public:
   using OpConversionPattern<ConstSizeOp>::OpConversionPattern;
 

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

 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
-static RankedTensorType getExtentTensorType(OpBuilder &builder) {
-  return RankedTensorType::get({ShapedType::kDynamicSize},
-                               builder.getIndexType());
+static RankedTensorType getExtentTensorType(MLIRContext *ctx) {
+  return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
 }
 
 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 (type.isa<ShapeType>() || type == getExtentTensorType(builder))
+  if (type.isa<ShapeType>() ||
+      type == getExtentTensorType(builder.getContext()))
     return builder.create<ConstShapeOp>(loc, type,
                                         value.cast<DenseIntElementsAttr>());
   if (type.isa<SizeType>())
     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
   if (type.isa<WitnessType>())
     return builder.create<ConstWitnessOp>(loc, type, value.cast<BoolAttr>());
   return nullptr;
 }

 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
   auto type = getOperand().getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return nullptr;
   Builder builder(getContext());
   return builder.getIndexTensorAttr(type.getShape());
 }
 
+static LogicalResult verify(ShapeOfOp op) {
+  Type argTy = op.arg().getType();
+  Type resultTy = op.result().getType();
+  if (argTy.isa<ValueShapeType>()) {
+    if (!resultTy.isa<ShapeType>())
+      return op.emitOpError()
+             << "if operand is of type `value_shape` then the result must be "
+                "of type `shape` to propagate potential error shapes";
+  } else {
+    assert(argTy.isa<ShapedType>());
+    if (resultTy != getExtentTensorType(op.getContext()))
+      return op.emitOpError() << "if operand is a shaped type then the result "
+                                 "must be an extent tensor";
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // SizeToIndexOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

func @shape_of_unranked(%arg : tensor<*xf32>) {
   // CHECK-DAG: %[[C1:.*]] = constant 1 : index
   // CHECK:     scf.for %[[I:.*]] = %[[C0]] to %[[RANK]] step %[[C1]] {
   // CHECK-DAG:   %[[DIM:.]] = dim %[[ARG]], %[[I]] : tensor<*xf32>
   // CHECK-DAG:   %[[DIM_INT:.*]] = index_cast %[[DIM]] : index to i64
   // CHECK-DAG:   store %[[DIM_INT]], %[[SHAPE_MEM]][%[[I]]] : memref<?xi64>
   // CHECK:     }
   // CHECK-DAG: %[[SHAPE_INT:.*]] = tensor_load %[[SHAPE_MEM]] : memref<?xi64>
   // CHECK-DAG: %[[SHAPE:.*]] = index_cast %[[SHAPE_INT]] : tensor<?xi64> to tensor<?xindex>
-  %shape = shape.shape_of %arg : tensor<*xf32>
+  %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
   return
 }
 
 // -----
 
 // CHECK-LABEL:  @shape_eq
 // CHECK-SAME:   (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>) -> i1
 func @shape_eq(%a : tensor<?xindex>, %b : tensor<?xindex>) -> i1 {

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 
 // Lower `shape_of` for statically shaped tensor.
 // CHECK-LABEL: @shape_of_stat
 // CHECK-SAME: (%[[ARG:.*]]: tensor<1x2x3xf32>)
 func @shape_of_stat(%arg : tensor<1x2x3xf32>) {
   // CHECK-DAG: %[[C1:.*]] = constant 1 : index
   // CHECK-DAG: %[[C2:.*]] = constant 2 : index
   // CHECK-DAG: %[[C3:.*]] = constant 3 : index
-  // CHECK-DAG: %[[SHAPE:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]]) : tensor<3xindex>
-  %shape = shape.shape_of %arg : tensor<1x2x3xf32>
+  // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]]) : tensor<3xindex>
+  // CHECK-DAG: %[[SHAPE:.*]] = tensor_cast %[[SHAPE_UNCASTED]] : tensor<3xindex> to tensor<?xindex>
+  %shape = shape.shape_of %arg : tensor<1x2x3xf32> -> tensor<?xindex>
   return
 }
 
 // -----
 
 // Lower `shape_of` for dynamically shaped tensor.
 // CHECK-LABEL: @shape_of_dyn
 // CHECK-SAME: (%[[ARG:.*]]: tensor<1x5x?xf32>)
 func @shape_of_dyn(%arg : tensor<1x5x?xf32>) {
   // CHECK-DAG: %[[C1:.*]] = constant 1 : index
   // CHECK-DAG: %[[C5:.*]] = constant 5 : index
   // CHECK-DAG: %[[C2:.*]] = constant 2 : index
   // CHECK-DAG: %[[DYN_DIM:.*]] = dim %[[ARG]], %[[C2]] : tensor<1x5x?xf32>
-  // CHECK-DAG: %[[SHAPE:.*]] = tensor_from_elements(%[[C1]], %[[C5]], %[[DYN_DIM]]) : tensor<3xindex>
-  %shape = shape.shape_of %arg : tensor<1x5x?xf32>
+  // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements(%[[C1]], %[[C5]], %[[DYN_DIM]]) : tensor<3xindex>
+  // CHECK-DAG: %[[SHAPE:.*]] = tensor_cast %[[SHAPE_UNCASTED]] : tensor<3xindex> to tensor<?xindex>
+  %shape = shape.shape_of %arg : tensor<1x5x?xf32> -> tensor<?xindex>
   return
 }
 
 // -----
 
 // Convert `rank` to `dim` of the first dimension.
 // CHECK-LABEL: @rank
 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index

 // Express `get_extent` as `std.dim` when it relies directly on the outcome of a
 // `shape_of` operation.
 // CHECK-LABEL: @get_extent_shape_of
 // CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
 func @get_extent_shape_of(%arg : tensor<2x3xf32>, %idx : !shape.size)
     -> !shape.size {
   // CHECK: %[[RESULT:.*]] = dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
   // CHECK: return %[[RESULT]] : index
-  %shape = shape.shape_of %arg : tensor<2x3xf32>
-  %result = shape.get_extent %shape, %idx : !shape.shape
+  %shape = shape.shape_of %arg : tensor<2x3xf32> -> tensor<?xindex>
+  %result = shape.get_extent %shape, %idx : tensor<?xindex>
   return %result : !shape.size
 }
 
 // -----
 
 // Express `get_extent` as `std.extract_element` when it relies directly on the
 // outcome of a `from_extent_tensor` operation.
 // CHECK-LABEL: @get_extent_from_extent_tensor

mlir/test/Dialect/Shape/canonicalize.mlir

-// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize <%s | FileCheck %s
+// 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] : !shape.shape
-  %0 = "shape.shape_of"(%arg0) : (tensor<2x3x4xf32>) -> !shape.shape
-  return %0 : !shape.shape
+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>
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK: shape.const_shape [2, 3] : !shape.shape

 
 // -----
 
 // Canonicalize `rank` when shape is derived from ranked tensor.
 // CHECK-LABEL: @canonicalize_rank
 func @canonicalize_rank(%arg : tensor<1x2x?xf32>) -> !shape.size {
   // CHECK-DAG: %[[RESULT:.*]] = shape.const_size 3
   // CHECK-DAG: return %[[RESULT]] : !shape.size
-  %shape = shape.shape_of %arg : tensor<1x2x?xf32>
-  %rank = shape.rank %shape : !shape.shape
+  %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> tensor<?xindex>
+  %rank = shape.rank %shape : tensor<?xindex>
   return %rank : !shape.size
 }
 
 // -----
 
 // Do not canonicalize `rank` when shape is derived from unranked tensor.
 // CHECK-LABEL: @dont_canonicalize_rank
 // CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>) -> !shape.size
 func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> !shape.size {
-  // CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32>
+  // CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32> -> tensor<?xindex>
   // CHECK-DAG: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
   // CHECK-DAG: return %[[SIZE]] : !shape.size
-  %shape = shape.shape_of %arg : tensor<*xf32>
-  %rank = shape.rank %shape : !shape.shape
+  %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
+  %rank = shape.rank %shape : tensor<?xindex>
   return %rank : !shape.size
 }
 
 // Canonicalize redundant conversion from `index` to `size` and back.
 // CHECK-LABEL: @index_to_size_to_index
 // CHECK-SAME: (%[[IDX:.*]]: index) -> index
 func @index_to_size_to_index(%index : index) -> index {
   // CHECK: return %[[IDX]] : index

 
 // 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 [] : !shape.shape
-  %1 = shape.shape_of %arg0 : tensor<?xf32>
-  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
+  %1 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Do not canonicalize cstr_broadcastable checks with 2 unknowns
 // CHECK-LABEL: @cstr_broadcastable_unknown
 func @cstr_broadcastable_unknown(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>) {
   // CHECK-NEXT: shape.shape_of %arg0
   // CHECK-NEXT: shape.shape_of %arg1
   // CHECK-NEXT: shape.cstr_broadcastable
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.shape_of %arg0 : tensor<?xf32>
-  %1 = shape.shape_of %arg1 : tensor<?xf32>
-  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
+  %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
+  %1 = shape.shape_of %arg1 : tensor<?xf32> -> tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Scalars are safe to broadcast to unranked sizes.
 // CHECK-LABEL: @cstr_broadcastable_scalar_unranked
 func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>, %arg1 : tensor<index>) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
-  %0 = shape.shape_of %arg1 : tensor<index>
-  %1 = shape.shape_of %arg0 : tensor<*xf32>
-  %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
+  %0 = shape.shape_of %arg1 : tensor<index> -> tensor<?xindex>
+  %1 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
+  %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // Fold `shape_eq` for equal and constant shapes.
 // CHECK-LABEL: @shape_eq_fold_1

mlir/test/Dialect/Shape/invalid.mlir

 
 // -----
 
 func @assuming_all_op_too_few_operands() {
   // expected-error@+1 {{no operands specified}}
   %w0 = shape.assuming_all
   return
 }
+
+// -----
+
+func @shape_of(%value_arg : !shape.value_shape,
+               %shaped_arg : tensor<?x3x4xf32>) {
+  // expected-error@+1 {{if operand is of type `value_shape` then the result must be of type `shape` to propagate potential error shapes}}
+  %0 = shape.shape_of %value_arg : !shape.value_shape -> tensor<?xindex>
+}
+
+// -----
+
+func @shape_of(%value_arg : !shape.value_shape,
+               %shaped_arg : tensor<?x3x4xf32>) {
+  // expected-error@+1 {{if operand is a shaped type then the result must be an extent tensor}}
+  %1 = shape.shape_of %shaped_arg : tensor<?x3x4xf32> -> !shape.shape
+}
+

mlir/test/Dialect/Shape/ops.mlir

 
 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>
   return
 }
 
-func @test_shape_of(%arg0: tensor<?xf32>) -> !shape.shape {
-  %0 = shape.shape_of %arg0 : tensor<?xf32>
-  return %0 : !shape.shape
+func @test_shape_of(%arg0: tensor<?xf32>) -> tensor<?xindex> {
+  %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
+  return %0 : tensor<?xindex>
 }
 
 func @test_constraints() {
   %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