[MLIR][Shape] Canonicalize casted dynamic extent tensor

[frgossen]

Summary

Depends On D99159

[tpopp]

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

1081

Can we generalize this? Like can we also canonicalize

%0 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<?xindex>
%1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>

I would assume it can work both ways. The only case I'm not sure is safe is with unranked tensors like.

%0 = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
%1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>

[herhut]

This moves us a little further in the direction of having more precise types in the shape dialect in the non-error bearing case. I think this is the right direction to go but also want to give Jacques a chance to comment.

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

1029

This change seems unrelated. Can you add a motivation to the diff description?

1031

nit: Move into conditional.

[herhut]

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

1081

Good point, yes, the former would also be a valid canonicalization. The latter is not, unless we give shape_of asserting behavior wrt. its return type, which tensor.cast has. I think we should not and rather use the cast in this case.

This is independent of the fact that asserting behavior is not implemented for tensor.cast I think.

[frgossen]

Address comments

[frgossen]

Split CL

[frgossen]

Thanks!

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

1029

Right, it does not really belong here. Will make a separate CL

[frgossen]

Address comments

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

 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>();
+  Type type;
+  Type indexTy = builder.getIndexType();
+  if (auto shapedTy = arg.getType().dyn_cast<ShapedType>()) {
+    type = RankedTensorType::get(
+        {shapedTy.hasRank() ? shapedTy.getRank() : ShapedType::kDynamicSize},
+        indexTy);
+  } else {
+    type = builder.getType<ShapeType>();
+  }
   return ShapeOfOp::build(builder, result, type, arg);
 }
 
 namespace {
 struct ShapeOfWithTensor : public OpRewritePattern<shape::ShapeOfOp> {
   using OpRewritePattern<shape::ShapeOfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(shape::ShapeOfOp op,
                                 PatternRewriter &rewriter) const override {
     if (!op.arg().getType().isa<ShapedType>())
       return failure();
     if (op.getType().isa<ShapedType>())
       return failure();
 
     rewriter.replaceOpWithNewOp<shape::ShapeOfOp>(op.getOperation(), op.arg());
     return success();
   }
 };
+
+// Canonicalize
+// ```
+// %0 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<3xindex>
+// %1 = tensor.cast %0 : tensor<3xindex> to tensor<?xindex>
+// ```
+// to
+// ```
+// %1 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<?xindex>
+// ```
+struct ShapeOfCastedToDynamicExtentTensor
+    : public OpRewritePattern<tensor::CastOp> {
+  using OpRewritePattern<tensor::CastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::CastOp op,
+                                PatternRewriter &rewriter) const override {
+    Type extentTensorTy = getExtentTensorType(op.getContext());
+    if (op.getType() != extentTensorTy)
+      return failure();
+
+    auto shapeOfOp = op.source().getDefiningOp<ShapeOfOp>();
+    if (!shapeOfOp)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<ShapeOfOp>(op, extentTensorTy, shapeOfOp.arg());
+    return success();
+  }
+};
 } // namespace
 
 void ShapeOfOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
                                             MLIRContext *context) {
-  patterns.insert<ShapeOfWithTensor>(context);
+  patterns.insert<ShapeOfCastedToDynamicExtentTensor, ShapeOfWithTensor>(
+      context);
 }
 
 //===----------------------------------------------------------------------===//
 // SizeToIndexOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as

mlir/test/Dialect/Shape/canonicalize.mlir

%results:3 = shape.assuming %witness
       -> (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) {
     %inner = "some.tensor"() : () -> tensor<2x3xf32>
     shape.assuming_yield %inner, %outer, %arg
         : tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>
   }
   return %results#0, %results#1, %results#2
       : tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @casted_dynamic_extent_tensor
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?x?x?xf32>) -> tensor<?xindex>
+func @casted_dynamic_extent_tensor(%arg : tensor<?x?x?xf32>)
+    -> tensor<?xindex> {
+  // CHECK: %[[RESULT:.*]] = shape.shape_of %[[ARG]] : tensor<?x?x?xf32> -> tensor<?xindex>
+  // CHECK: return %[[RESULT]] : tensor<?xindex>
+  %0 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<3xindex>
+  %1 = tensor.cast %0 : tensor<3xindex> to tensor<?xindex>
+  return %1 : tensor<?xindex>
+}