[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

27

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

394

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

395

Doesn't it just return the shape?

396

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

27

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

394

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.

395

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

396

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

396

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

643

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-shape tuple or a shaped operand as an argument
+    and returns a description of the shape. Because this operation cannot yield
+    error shapes the result is an extent tensor `tensor<?xindex>`, not a shape
+    type.
+  }];
+
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$arg);
-  let results = (outs Shape_ShapeType:$result);
+  let results = (outs Shape_ExtentTensorType:$result);
 
   let assemblyFormat = "$arg `:` type($arg) attr-dict";
 
   let hasFolder = 1;
 }
 
 def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let summary = "Casts between index types of the shape and standard dialect";

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

 using namespace mlir;
 using namespace mlir::shape;
 
 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
 namespace {
-RankedTensorType getExtentTensorType(OpBuilder &builder) {
+RankedTensorType inline getExtentTensorType(OpBuilder &builder) {
   return RankedTensorType::get({ShapedType::kDynamicSize},
                                builder.getIndexType());
 }
 } // namespace
 
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<

 
 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])
     return arg;
   return {};
 }
 
 void SizeToIndexOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<IndexToSizeToIndexCanonicalization>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//

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>
+  // 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>
   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>
+  // 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>
   return
 }
 
 // -----
 
 // Convert `rank` to `dim` of the first dimension.
 // CHECK-LABEL: @rank

 // `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
+  %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>
+  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
+  %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: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
   // CHECK-DAG: return %[[SIZE]] : !shape.size
   %shape = shape.shape_of %arg : tensor<*xf32>
-  %rank = shape.rank %shape : !shape.shape
+  %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
+  %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
+  %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
+  %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }

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