}
};
+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:
BinaryOpConversion<AddOp, AddIOp>,
BinaryOpConversion<MulOp, MulIOp>,
ConstSizeOpConverter,
+ RankOpConverter,
ShapeOfOpConversion>(ctx);
// clang-format on
}
}
// CHECK: %[[C1:.*]] = constant 1 : index
// CHECK: return %[[C1]] : index
-
// -----
// Lower `shape_of` for statically shaped tensor.
%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
+}