[MLIR][Shape] Add canonicalization pattern for `shape.rank`

[frgossen]

Summary

Replace any rank(shape_of(tensor)) that relies on a ranked tensor with the
corresponding constant const_size.

Depends On D82076

[silvas]

Awesome, thanks!

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

461

nit: getDefiningOp<ShapeOfOp>()

[rriddle]

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

364

nit: inline this variable into its only use.

455

Can you documentation to this pattern? Even just a little code block of the transformation being performed.

[frgossen]

Address comments

[jpienaar]

Nice, thanks

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

461

Is there a : missing before tensor<...> ?

[rriddle]

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

454

Please use /// for top-level comments.

[frgossen]

Address comments

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

let description = [{
     ```mlir
     %x = shape.const_size 10
     ```
   }];
 
   let arguments = (ins IndexAttr:$value);
   let results = (outs Shape_SizeType:$result);
 
+  let builders = [
+    OpBuilder<"OpBuilder &builder, OperationState &result, int64_t value">
+  ];
+
   let assemblyFormat = "$value attr-dict";
   let hasFolder = 1;
 }
 
 def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
   let summary = "Creates a shape from extents";
   let description = [{
     Creates a shape from multiple SSA values representing the extents of
def Shape_GetSizeOp : Shape_Op<"get_size", [NoSideEffect]> {
   }];
 
   let arguments = (ins Shape_ShapeType:$shape);
   let results = (outs Shape_SizeType:$size);
 
   let assemblyFormat = "attr-dict $shape";
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let summary = "Creates a dimension tensor from a shape";
   let description = [{
     Converts a shape to a 1D integral tensor of extents. The number of elements
     in the tensor equals the rank of the shape, and the elements equal the
     extents of the shape.

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

OpFoldResult CstrEqOp::fold(ArrayRef<Attribute> operands) {
   // cannot if there are any non-const inputs.
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
+void ConstSizeOp::build(OpBuilder &builder, OperationState &result,
+                        int64_t value) {
+  auto attr = builder.getIndexAttr(value);
+  build(builder, result, attr);
+}
+
 OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
 
 void ConstSizeOp::getAsmResultNames(
     llvm::function_ref<void(Value, StringRef)> setNameFn) {
   SmallString<4> buffer;
   llvm::raw_svector_ostream os(buffer);
   os << "c" << value();
   setNameFn(getResult(), os.str());
OpFoldResult GetSizeOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (!shape)
     return {};
   int64_t size = shape.end() - shape.begin();
   Builder builder(getContext());
   return builder.getIndexAttr(size);
 }
 
+namespace {
+struct GetSizeShapeOfCanonicalizationPattern
+    : public OpRewritePattern<GetSizeOp> {
+  using OpRewritePattern<GetSizeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GetSizeOp op,
+                                PatternRewriter &rewriter) const override {
+    auto shapeOfOp = dyn_cast_or_null<ShapeOfOp>(op.shape().getDefiningOp());
+    if (!shapeOfOp)
+      return failure();
+    auto rankedTensorType =
+        shapeOfOp.arg().getType().dyn_cast<RankedTensorType>();
+    if (!rankedTensorType)
+      return failure();
+    int64_t shapeSize = rankedTensorType.getRank();
+    rewriter.replaceOpWithNewOp<ConstSizeOp>(op.getOperation(), shapeSize);
+    return success();
+  }
+};
+} // namespace
+
+void GetSizeOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                            MLIRContext *context) {
+  patterns.insert<GetSizeShapeOfCanonicalizationPattern>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // NumElementsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
 
   // Fold only when argument constant.
   Attribute shape = operands[0];

mlir/test/Dialect/Shape/canonicalize.mlir

 // CHECK-LABEL: @dont_fold_size
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size
 func @dont_fold_size(%shape : !shape.shape) -> !shape.size {
   // CHECK-DAG: %[[RESULT:.*]] = shape.get_size %[[SHAPE]]
   // CHECK-DAG: return %[[RESULT]] : !shape.size
   %size = shape.get_size %shape
   return %size : !shape.size
 }
+
+// -----
+
+// Canonicalize `get_size` when shape is derived from ranked tensor.
+// CHECK-LABEL: @canonicalize_size
+func @canonicalize_size(%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>
+%size = shape.get_size %shape
+return %size : !shape.size
+}
+
+// -----
+
+// Do not canonicalize `get_size` when shape is derived from unranked tensor.
+// CHECK-LABEL: @dont_canonicalize_size
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>) -> !shape.size
+func @dont_canonicalize_size(%arg : tensor<*xf32>) -> !shape.size {
+// CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32>
+// CHECK-DAG: %[[SIZE:.*]] = shape.get_size %[[SHAPE]]
+// CHECK-DAG: return %[[SIZE]] : !shape.size
+%shape = shape.shape_of %arg : tensor<*xf32>
+%size = shape.get_size %shape
+return %size : !shape.size
+}