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

[frgossen]

Summary

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.

[frgossen]

Rebase

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

public:
   matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 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();
 }
 
 namespace {
 /// Conversion pass.
 class ConvertShapeToStandardPass

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

func @get_extent(%shape : tensor<?xindex>, %idx : !shape.size) -> !shape.size {
   // CHECK: shape.get_extent
   %result = shape.get_extent %shape, %idx
       : tensor<?xindex>, !shape.size -> !shape.size
   return %result : !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
 // CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
 func @get_extent_shape_of(%arg : tensor<2x3xf32>, %idx : index) -> index {
   // CHECK: %[[RESULT:.*]] = dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
   // CHECK: return %[[RESULT]] : index
   %shape = shape.shape_of %arg : tensor<2x3xf32> -> tensor<?xindex>