[MLIR][Shape] Limit shape to standard lowerings to their supported types

[frgossen]

Summary

The lowering does not support all types for its source operations. This change
makes the patterns fail in a well-defined manner.

Depends On D84442

[rriddle]

Are there any test cases for these new failure cases? If not could you please add some?

[frgossen]

Add tests

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

 template <typename SrcOpTy, typename DstOpTy>
 class BinaryOpConversion : public OpConversionPattern<SrcOpTy> {
 public:
   using OpConversionPattern<SrcOpTy>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    typename SrcOpTy::Adaptor adaptor(operands);
-    rewriter.replaceOpWithNewOp<DstOpTy>(op, adaptor.lhs(), adaptor.rhs());
+    typename SrcOpTy::Adaptor transformed(operands);
+
+    // For now, only error-free types are supported by this lowering.
+    if (op.getType().template isa<SizeType>())
+      return failure();
+
+    rewriter.replaceOpWithNewOp<DstOpTy>(op, transformed.lhs(),
+                                         transformed.rhs());
     return success();
   }
 };
 } // namespace
 
 namespace {
 class ConstSizeOpConversion : public OpConversionPattern<ConstSizeOp> {
 public:
public:
   matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 LogicalResult ShapeOfOpConversion::matchAndRewrite(
     ShapeOfOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
-  ShapeOfOp::Adaptor transformed(operands);
-  auto loc = op.getLoc();
-  auto tensorVal = transformed.arg();
-  auto tensorTy = tensorVal.getType();
+
+  // For now, only error-free types are supported by this lowering.
+  if (op.getType().isa<ShapeType>())
+    return failure();
 
   // For unranked tensors `shape_of` lowers to `scf` and the pattern can be
   // found in the corresponding pass.
+  ShapeOfOp::Adaptor transformed(operands);
+  Value tensorVal = transformed.arg();
+  Type tensorTy = tensorVal.getType();
   if (tensorTy.isa<UnrankedTensorType>())
     return failure();
 
   // Build values for individual dimensions.
   SmallVector<Value, 8> dimValues;
-  auto rankedTensorTy = tensorTy.cast<RankedTensorType>();
+  RankedTensorType rankedTensorTy = tensorTy.cast<RankedTensorType>();
   int64_t rank = rankedTensorTy.getRank();
+  auto loc = op.getLoc();
   for (int64_t i = 0; i < rank; i++) {
     if (rankedTensorTy.isDynamicDim(i)) {
-      auto dimVal = rewriter.create<DimOp>(loc, tensorVal, i);
+      Value dimVal = rewriter.create<DimOp>(loc, tensorVal, i);
       dimValues.push_back(dimVal);
     } else {
       int64_t dim = rankedTensorTy.getDimSize(i);
-      auto dimVal = rewriter.create<ConstantIndexOp>(loc, dim);
+      Value dimVal = rewriter.create<ConstantIndexOp>(loc, dim);
       dimValues.push_back(dimVal);
     }
   }
 
   // Materialize extent tensor.
   Value staticExtentTensor =
       rewriter.create<TensorFromElementsOp>(loc, dimValues);
   rewriter.replaceOpWithNewOp<TensorCastOp>(op, staticExtentTensor,

 };
 } // namespace
 
 LogicalResult GetExtentOpConverter::matchAndRewrite(
     GetExtentOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   GetExtentOp::Adaptor transformed(operands);
 
-  // Derive shape extent directly from shape origin if possible.
-  // This circumvents the necessity to materialize the shape in memory.
+  // For now, only error-free types are supported by this lowering.
+  if (op.getType().isa<SizeType>())
+    return failure();
+
+  // Derive shape extent directly from shape origin if possible. This
+  // circumvents the necessity to materialize the shape in memory.
   if (auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>()) {
-    rewriter.replaceOpWithNewOp<DimOp>(op, shapeOfOp.arg(), transformed.dim());
+    if (shapeOfOp.arg().getType().isa<ShapedType>()) {
+      rewriter.replaceOpWithNewOp<DimOp>(op, shapeOfOp.arg(),
+                                         transformed.dim());
       return success();
     }
+  }
 
   rewriter.replaceOpWithNewOp<ExtractElementOp>(op, rewriter.getIndexType(),
                                                 transformed.shape(),
                                                 ValueRange{transformed.dim()});
   return success();
 }
 
 namespace {
void ConvertShapeToStandardPass::runOnOperation() {
   target.addLegalOp<FuncOp, ModuleOp, ModuleTerminatorOp>();
 
   // Setup conversion patterns.
   OwningRewritePatternList patterns;
   populateShapeToStandardConversionPatterns(patterns, &ctx);
 
   // 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

 // RUN: mlir-opt --split-input-file --convert-shape-to-std --verify-diagnostics %s | FileCheck %s
 
 // Lower binary ops.
 // CHECK-LABEL: @binary_ops
 // CHECK-SAME: (%[[LHS:.*]]: index, %[[RHS:.*]]: index)
 func @binary_ops(%lhs : index, %rhs : index) {
   // CHECK: addi %[[LHS]], %[[RHS]] : index
   %sum = shape.add %lhs, %rhs : index, index -> index
   // CHECK: muli %[[LHS]], %[[RHS]] : index
   %product = shape.mul %lhs, %rhs : index, index -> index
   return
 }
 
 // -----
 
+// Don't lower binary ops when they operate on `shape.size`.
+// CHECK-LABEL: @binary_ops_on_size
+// CHECK-SAME: (%[[LHS:.*]]: !shape.size, %[[RHS:.*]]: !shape.size)
+func @binary_ops_on_size(%lhs : !shape.size, %rhs : !shape.size) {
+  // CHECK: shape.add %[[LHS]], %[[RHS]] : !shape.size, !shape.size -> !shape.size
+  // CHECK: shape.mul %[[LHS]], %[[RHS]] : !shape.size, !shape.size -> !shape.size
+  %sum = shape.add %lhs, %rhs : !shape.size, !shape.size -> !shape.size
+  %prod = shape.mul %lhs, %rhs : !shape.size, !shape.size -> !shape.size
+  return
+}
+
+// -----
+
+// Don't lower `shape_of` with `shape.shape` type.
+// CHECK-LABEL: @shape_of
+// CHECK-SAME: (%[[ARG:.*]]: tensor<1x2x3xf32>)
+func @shape_of_stat(%arg : tensor<1x2x3xf32>) {
+  // CHECK: shape.shape_of %[[ARG]] : tensor<1x2x3xf32> -> !shape.shape
+  %shape = shape.shape_of %arg : tensor<1x2x3xf32> -> !shape.shape
+  return
+}
+
+// -----
+
 // Lower `shape_of` for statically shaped tensor.
 // CHECK-LABEL: @shape_of_stat
 // CHECK-SAME: (%[[ARG:.*]]: tensor<1x2x3xf32>)
 func @shape_of_stat(%arg : tensor<1x2x3xf32>) {
   // CHECK-DAG: %[[C1:.*]] = constant 1 : index
   // CHECK-DAG: %[[C2:.*]] = constant 2 : index
   // CHECK-DAG: %[[C3:.*]] = constant 3 : index
   // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements(%[[C1]], %[[C2]], %[[C3]]) : tensor<3xindex>
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 `get_extent` if it is of type `shape.size`.
+// CHECK-LABEL: @get_extent
+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
+}
+
+// -----
+
 // 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>