[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

465

nit: getDefiningOp<ShapeOfOp>()

[rriddle]

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

369

nit: inline this variable into its only use.

459

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

465

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

[rriddle]

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

458

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_RankOp : Shape_Op<"rank", [NoSideEffect]> {
   }];
 
   let arguments = (ins Shape_ShapeType:$shape);
   let results = (outs Shape_SizeType:$rank);
 
   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) {
+  build(builder, result, builder.getIndexAttr(value));
+}
+
 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 RankOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (!shape)
     return {};
   int64_t rank = shape.getNumElements();
   Builder builder(getContext());
   return builder.getIndexAttr(rank);
 }
 
+/// Evaluate the `rank` operation for shapes of ranked tensors at compile time.
+/// Constant folding fails in cases where only the rank is constant, not the
+/// shape itself.
+/// This canonicalization matches `shape.rank(shape.shape_of(%ranked_tensor))`.
+///
+/// Example:
+///
+/// %shape = shape.shape_of %ranked_tensor : tensor<1x2x?xf32>
+/// %rank = shape.rank %shape
+///
+/// becomes
+///
+/// %rank = shape.const_size 3
+
+namespace {
+struct RankShapeOfCanonicalizationPattern : public OpRewritePattern<RankOp> {
+  using OpRewritePattern<RankOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(RankOp op,
+                                PatternRewriter &rewriter) const override {
+    auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>();
+    if (!shapeOfOp)
+      return failure();
+    auto rankedTensorType =
+        shapeOfOp.arg().getType().dyn_cast<RankedTensorType>();
+    if (!rankedTensorType)
+      return failure();
+    int64_t rank = rankedTensorType.getRank();
+    rewriter.replaceOpWithNewOp<ConstSizeOp>(op.getOperation(), rank);
+    return success();
+  }
+};
+} // namespace
+
+void RankOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                         MLIRContext *context) {
+  patterns.insert<RankShapeOfCanonicalizationPattern>(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_rank
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size
 func @dont_fold_rank(%shape : !shape.shape) -> !shape.size {
   // CHECK-DAG: %[[RESULT:.*]] = shape.rank %[[SHAPE]]
   // CHECK-DAG: return %[[RESULT]] : !shape.size
   %rank = shape.rank %shape
   return %rank : !shape.size
 }
+
+// -----
+
+// 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
+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
+return %rank : !shape.size
+}