diff --git a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp --- a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp +++ b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp @@ -231,6 +231,10 @@ LogicalResult RankOpConverter::matchAndRewrite(shape::RankOp op, ArrayRef operands, ConversionPatternRewriter &rewriter) const { + // For now, this lowering supports only error-free types. + if (op.getType().isa()) + return failure(); + shape::RankOp::Adaptor transformed(operands); rewriter.replaceOpWithNewOp(op, transformed.shape(), 0); return success(); diff --git a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir --- a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir +++ b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir @@ -90,6 +90,16 @@ // ----- +// Don't lower `rank` if type is not error-free. +// CHECK-LABEL: @rank +func @rank(%shape : !shape.shape) { + // CHECK: shape.rank + %rank = shape.rank %shape : !shape.shape -> !shape.size + return +} + +// ----- + // Express `get_extent` as `std.dim` when it relies directly on the outcome of a // `shape_of` operation. // CHECK-LABEL: @get_extent_shape_of