[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 {
 /// Type conversions.
 class ShapeTypeConverter : public TypeConverter {
void ConvertShapeToStandardPass::runOnOperation() {
 
   // Setup conversion patterns.
   OwningRewritePatternList patterns;
   populateShapeToStandardConversionPatterns(patterns, &ctx);
   populateFuncOpTypeConversionPattern(patterns, &ctx, typeConverter);
 
   // Apply conversion.
   auto module = getOperation();
-  if (failed(applyFullConversion(module, target, patterns)))
+  if (failed(applyPartialConversion(module, target, patterns)))
     signalPassFailure();
 }
 
 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
       AnyOpConversion,

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

func @rank(%shape : tensor<?xindex>) -> index {
   // CHECK: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
   // CHECK: return %[[RESULT]] : index
   %rank = shape.rank %shape : tensor<?xindex> -> index
   return %rank : index
 }
 
 // -----
 
+// 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>