[MLIR][Shape] Limit `shape.rank` lowering to its extent tensor variant

When lowering to the standard dialect, we currently support only the extent
tensor variant of the shape.rank operation. This change lets the conversion
pattern fail in a well-defined manner.

Differential Revision: https://reviews.llvm.org/D84852

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

@@ -231,6 +231,10 @@ public:
 LogicalResult
 RankOpConverter::matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
                                  ConversionPatternRewriter &rewriter) const {
+  // For now, this lowering supports only error-free types.
+  if (op.getType().isa<SizeType>())
+    return failure();
+
   shape::RankOp::Adaptor transformed(operands);
   rewriter.replaceOpWithNewOp<DimOp>(op, transformed.shape(), 0);
   return success();

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

@@ -90,6 +90,16 @@ func @get_extent(%shape : tensor<?xindex>, %idx : !shape.size) -> !shape.size {
 
 // -----
 
+// 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