[MLIR][Shape] Derive more concrete type for `shape.shape_of`

[frgossen]

Summary

Also create all extent tensor constants with const_shape op.

[jpienaar]

This seems as expected (capturing what is known in the non-error case when building), I'm surprised no tests needed updating - guessing as we don't have one exercising this builder. Could you perhaps add a small C++ unit test for it?

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

1154–1155

Could you try one where you do

if (auto shapedTy = ...) {
  ...
  return ..
}
return ShapeOfOp::build(builder, result, builder.getType<ShapeType>(), arg);

? Seems like it could be simpler

1155

Nit: reduce scope by moving to just before use

1158

Hoist this out (I think it will both reduce whitespace and read easier)

[herhut]

LGTM for the direction this takes. Please address nits.

[frgossen]

Address comments

[tpopp]

Is this no longer needed?

[frgossen]

It is but there are some internal test failing that I would like to fix before landing this

[frgossen]

PTAL

[frgossen]

Rebase

mlir/include/mlir/Dialect/Shape/IR/Shape.h

 namespace mlir {
 class PatternRewriter;
 
 namespace shape {
 
 /// Alias type for extent tensors.
 RankedTensorType getExtentTensorType(MLIRContext *ctx);
 
+// Check if a type is an extent tensor, e.g., tensor<?xindex>.
+bool isExtentTensorType(Type);
+
 // Given an input shape Value, try to obtain the shape's values.
 LogicalResult getShapeVec(Value input, SmallVectorImpl<int64_t> &shapeValues);
 
 /// The shape descriptor type represents rank and dimension sizes.
 class ShapeType : public Type::TypeBase<ShapeType, Type, TypeStorage> {
 public:
   using Base::Base;
 };

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

 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
 RankedTensorType shape::getExtentTensorType(MLIRContext *ctx) {
   return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
 }
 
+bool shape::isExtentTensorType(Type type) {
+  auto ranked = type.dyn_cast<RankedTensorType>();
+  return ranked && ranked.getRank() == 1 && ranked.getElementType().isIndex();
+}
+
 LogicalResult shape::getShapeVec(Value input,
                                  SmallVectorImpl<int64_t> &shapeValues) {
   if (auto inputOp = input.getDefiningOp<ShapeOfOp>()) {
     auto type = inputOp.arg().getType().dyn_cast<ShapedType>();
     if (!type.hasRank())
       return failure();
     shapeValues = llvm::to_vector<6>(type.getShape());
     return success();
#include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
   // 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.getContext()))
+  if (type.isa<ShapeType>() || isExtentTensorType(type))
     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>());
   if (ConstantOp::isBuildableWith(value, type))
     return builder.create<ConstantOp>(loc, type, value);

 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());
 }
 
 void ShapeOfOp::build(OpBuilder &builder, OperationState &result, Value arg) {
-  Type type = arg.getType().isa<ShapedType>()
-                  ? (Type)getExtentTensorType(builder.getContext())
-                  : (Type)builder.getType<ShapeType>();
-  return ShapeOfOp::build(builder, result, type, arg);
+  if (auto shapedTy = arg.getType().dyn_cast<ShapedType>()) {
+    int64_t rank =
+        shapedTy.hasRank() ? shapedTy.getRank() : ShapedType::kDynamicSize;
+    Type indexTy = builder.getIndexType();
+    Type extentTensorTy = RankedTensorType::get({rank}, indexTy);
+    return ShapeOfOp::build(builder, result, extentTensorTy, arg);
+  }
+  Type shapeTy = builder.getType<ShapeType>();
+  return ShapeOfOp::build(builder, result, shapeTy, arg);
 }
 
 namespace {
 struct ShapeOfWithTensor : public OpRewritePattern<shape::ShapeOfOp> {
   using OpRewritePattern<shape::ShapeOfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(shape::ShapeOfOp op,
                                 PatternRewriter &rewriter) const override {

mlir/test/Dialect/Shape/canonicalize.mlir

 }
 
 // -----
 
 // Partial folding.
 // CHECK-LABEL: @partial_folding
 // CHECK-SAME:  (%[[ARG:.*]]: !shape.shape)
 func @partial_folding(%arg0 : !shape.shape) -> !shape.shape {
-  // CHECK: %[[CST_SHAPE:.*]] = constant dense<[1, 2, 3]> : tensor<3xindex>
+  // CHECK: %[[CST_SHAPE:.*]] = shape.const_shape [1, 2, 3] : tensor<3xindex>
   // CHECK: %[[RESULT:.*]] = shape.broadcast %[[ARG]], %[[CST_SHAPE]] : !shape.shape, tensor<3xindex> -> !shape.shape
   // CHECK: return %[[RESULT]]
   %0 = shape.const_shape [2, 1] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : tensor<3xindex>
   %2 = shape.broadcast %0, %arg0, %1, %0 : !shape.shape, !shape.shape,
       tensor<3xindex>, !shape.shape -> !shape.shape
   return %2 : !shape.shape
 }
func @f() -> !shape.shape {
   return %0 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> tensor<2xindex> {
-  // CHECK: constant dense<[0, 1]> : tensor<2xindex>
+  // CHECK: shape.const_shape [0, 1] : tensor<2xindex>
   %cs = shape.const_shape [0, 1] : !shape.shape
   %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // -----
 
 // Basic case.
func @fold_assuming_all_single_element(%arg: tensor<?xindex>) {
   return
 }
 
 // -----
 
 // 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: constant dense<2> : 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>
   return %1 : tensor<1xindex>
 }
 
 // -----
 

 
 // CHECK-LABEL: min_same_arg
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)
 func @min_same_arg(%a: !shape.shape) -> !shape.shape {
   %1 = shape.min %a, %a : !shape.shape, !shape.shape -> !shape.shape
   // CHECK: return %[[SHAPE]]
   return %1 : !shape.shape
 }
-
 // ----
 
 // CHECK-LABEL: @cstr_broadcastable_folding
 func @cstr_broadcastable_folding(%arg : tensor<?x4xf32>) {
   // CHECK: const_witness true
   %0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex>
-  %1 = constant dense<[4]> : tensor<1xindex>
+  %1 = shape.const_shape [4] : tensor<1xindex>
   %2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex>
   "use"(%2) : (!shape.witness) -> ()
 }