[MLIR][Shape] Lower `shape.rank`

[frgossen]

Summary

Lower shape.rank to standard dialect.
A shape's size is the same as the extent of the first and only dimension of the
tensor<?xindex> it is represented by.

Depends On D82077

[frgossen]

Rebase

[frgossen]

Rebase

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

public:
   matchAndRewrite(ConstSizeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<ConstantIndexOp>(op.getOperation(),
                                                  op.value().getSExtValue());
     return success();
   }
 };
 
+class RankOpConverter : public OpConversionPattern<shape::RankOp> {
+public:
+  using OpConversionPattern<shape::RankOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    shape::RankOp::Adaptor transformed(operands);
+    rewriter.replaceOpWithNewOp<DimOp>(op.getOperation(), transformed.shape(),
+                                       0);
+    return success();
+  }
+};
+
 /// Type conversions.
 class ShapeTypeConverter : public TypeConverter {
 public:
   using TypeConverter::convertType;
 
   ShapeTypeConverter(MLIRContext *ctx) {
     // Add default pass-through conversion.
     addConversion([&](Type type) { return type; });

 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   populateWithGenerated(ctx, &patterns);
   // clang-format off
   patterns.insert<
       BinaryOpConversion<AddOp, AddIOp>,
       BinaryOpConversion<MulOp, MulIOp>,
       ConstSizeOpConverter,
+      RankOpConverter,
       ShapeOfOpConversion>(ctx);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // Convert `const_size` to `constant` op.
 // CHECK-LABEL: @size_const
 func @size_const() -> !shape.size {
   %c1 = shape.const_size 1
   return %c1 : !shape.size
 }
 // CHECK: %[[C1:.*]] = constant 1 : index
 // CHECK: return %[[C1]] : index
-
 // -----
 
 // 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
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>
   %shape = shape.shape_of %arg : tensor<1x5x?xf32>
   return
 }
+
+// -----
+
+// Convert `rank` to `dim` of the first dimension.
+// CHECK-LABEL: @rank
+// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
+func @rank(%shape : !shape.shape) -> !shape.size {
+  // CHECK-DAG: %[[C0:.*]] = constant 0 : index
+  // CHECK-DAG: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
+  // CHECK-DAG: return %[[RESULT]] : index
+  %rank = shape.rank %shape
+  return %rank : !shape.size
+}