[MLIR][Shape] Refactor verification

[frgossen]

Summary

Eliminate duplicate code for verification purposes. This concerns operations
that operate on error (size, shape, ...) or non-error values (index,
tensor<?xindex>).

Depends On D84438

[jpienaar]

This one was a bit restrictive and let to requiring some dynamic shapes where we had exact ranks, we can revise later, but that made canonicalization etc more tricky to reason about.

[frgossen]

Loosen verification

[jpienaar]

Could you sync to head as I think this has some overlap with change I landed recently based on this :)

[frgossen]

See https://github.com/llvm/llvm-project/commit/5142448a5e2aeeffefb3aabdb48f19033025bc09

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_RankOp : Shape_Op<"rank", [NoSideEffect]> {
 
   let arguments = (ins Shape_ShapeOrExtentTensorType:$shape);
   let results = (outs Shape_SizeOrIndexType:$rank);
 
   let assemblyFormat = "$shape `:` type($shape) `->` type($rank) attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
-  let verifier = [{ return ::verify(*this); }];
+  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
 }
 
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let summary = "Creates a dimension tensor from a shape";
   let description = [{
     Converts a shape to a 1D integral tensor of extents. The number of elements
     in the tensor equals the rank of the shape, and the elements equal the
     extents of the shape.
def Shape_GetExtentOp : Shape_Op<"get_extent", [NoSideEffect]> {
   ];
 
   let extraClassDeclaration = [{
     /// Get the `dim` value as integer if it is constant.
     Optional<int64_t> getConstantDim();
   }];
 
   let hasFolder = 1;
-  let verifier = [{ return ::verify(*this); }];
+  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
 }
 
 def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
   let summary = "Converts a standard index to a shape size";
   let description = [{
     Converts a standard index to a `shape.size`. This operation and its
     inverse, `size_to_index`, facilitate index conversion between the standard
     and the shape dialect.
def Shape_MulOp : Shape_Op<"mul", [Commutative, NoSideEffect]> {
 
   let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
   let results = (outs Shape_SizeOrIndexType:$result);
 
   let assemblyFormat = [{
     $lhs `,` $rhs `:` type($lhs) `,` type($rhs) `->` type($result) attr-dict
   }];
 
-  let verifier = [{ return ::verify(*this); }];
+  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
   let hasFolder = 1;
 }
 
 def Shape_NumElementsOp : Shape_Op<"num_elements", [NoSideEffect]> {
   let summary = "Returns the number of elements for a given shape";
   let description = [{
     Returns the number of elements for a given shape which is the product of its
     dimensions.
let description = [{
     returns a shape or extent tensor.
   }];
 
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$arg);
   let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
   let assemblyFormat = "$arg `:` type($arg) `->` type($result) attr-dict";
 
-  let verifier = [{ return ::verify(*this); }];
+  let verifier = [{ return ::verifyShapeOrExtentTensorOp(*this); }];
   let hasFolder = 1;
 }
 
 def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let summary = "Casts between index types of the shape and standard dialect";
   let description = [{
     Converts a `shape.size` to a standard index.
     This operation and its inverse, `index_to_size`, facilitate index conversion

mlir/lib/Dialect/Shape/IR/Shape.cpp

 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
 static RankedTensorType getExtentTensorType(MLIRContext *ctx) {
   return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
 }
 
-static bool isErrorPropagationPossible(ArrayRef<Type> operandTypes) {
+static bool isErrorPropagationPossible(TypeRange operandTypes) {
   for (Type ty : operandTypes)
     if (ty.isa<SizeType>() || ty.isa<ShapeType>() || ty.isa<ValueShapeType>())
       return true;
   return false;
 }
 
+static LogicalResult verifySizeOrIndexOp(Operation *op) {
+  assert(op != nullptr && op->getNumResults() == 1);
+  Type resultTy = op->getResultTypes().front();
+  if (isErrorPropagationPossible(op->getOperandTypes()) &&
+      !resultTy.isa<SizeType>())
+    return op->emitOpError()
+           << "if at least one of the operands can hold error values then the "
+              "result must be of type `size` to propagate them";
+  return success();
+}
+
+static LogicalResult verifyShapeOrExtentTensorOp(Operation *op) {
+  assert(op != nullptr && op->getNumResults() == 1);
+  Type resultTy = op->getResultTypes().front();
+  if (isErrorPropagationPossible(op->getOperandTypes()) &&
+      !resultTy.isa<ShapeType>())
+    return op->emitOpError()
+           << "if at least one of the operands can hold error values then the "
+              "result must be of type `shape` to propagate them";
+  return success();
+}
+
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
            WitnessType>();
OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   Builder builder(getContext());
   return builder.getIndexTensorAttr(extents);
 }
 
 //===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(GetExtentOp op) {
-  Type shapeTy = op.shape().getType();
-  Type dimTy = op.dim().getType();
-  Type extentTy = op.extent().getType();
-  if (isErrorPropagationPossible({shapeTy, dimTy})) {
-    if (!extentTy.isa<SizeType>())
-      op.emitError()
-          << "if at least one of the operands can hold error values then the "
-             "result must be of type `size` to propagate them";
-  } else {
-    if (extentTy.isa<SizeType>())
-      op.emitError() << "if none of the operands can hold error values then "
-                        "the result must be of type `index`";
-  }
-  return success();
-}
-
 Optional<int64_t> GetExtentOp::getConstantDim() {
   if (auto constSizeOp = dim().getDefiningOp<ConstSizeOp>())
     return constSizeOp.value().getLimitedValue();
   if (auto constantOp = dim().getDefiningOp<ConstantOp>())
     return constantOp.value().cast<IntegerAttr>().getInt();
   return llvm::None;
 }
 
if (shape.getType().isa<ShapeType>()) {
     build(builder, result, builder.getIndexType(), shape, dim);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(shape::RankOp op) {
-  if (op.shape().getType().isa<ShapeType>() &&
-      !op.rank().getType().isa<SizeType>())
-    return op.emitOpError()
-           << "if operand is of type `shape` then the result must be of type "
-              "`size` to propagate potential errors";
-  return success();
-}
-
 OpFoldResult shape::RankOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (!shape)
     return {};
   int64_t rank = shape.getNumElements();
   Builder builder(getContext());
   return builder.getIndexAttr(rank);
 }
OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
   Builder builder(getContext());
   return builder.getIndexAttr(product.getLimitedValue());
 }
 
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(MulOp op) {
-  Type resultTy = op.result().getType();
-  if (isErrorPropagationPossible({op.lhs().getType(), op.rhs().getType()})) {
-    if (!resultTy.isa<SizeType>())
-      return op.emitOpError()
-             << "if at least one of the operands can hold error values then "
-                "the result must be of type `size` to propagate them";
-  } else {
-    if (resultTy.isa<SizeType>())
-      return op.emitError() << "if none of the operands can hold error values "
-                               "then the result must be of type `index`";
-  }
-  return success();
-}
-
 OpFoldResult MulOp::fold(ArrayRef<Attribute> operands) {
   auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>();
   if (!lhs)
     return nullptr;
   auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>();
   if (!rhs)
     return nullptr;
   APInt folded = lhs.getValue() * rhs.getValue();
   Type indexTy = IndexType::get(getContext());
   return IntegerAttr::get(indexTy, folded);
 }
 
 //===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
   auto type = getOperand().getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return nullptr;
   Builder builder(getContext());
   return builder.getIndexTensorAttr(type.getShape());
 }
 
-static LogicalResult verify(ShapeOfOp op) {
-  Type resultTy = op.result().getType();
-  if (isErrorPropagationPossible(op.arg().getType())) {
-    if (!resultTy.isa<ShapeType>())
-      return op.emitOpError()
-             << "if operand is of type `value_shape` then the result must be "
-                "of type `shape` to propagate potential error shapes";
-  } else {
-    if (resultTy != getExtentTensorType(op.getContext()))
-      return op.emitOpError() << "if operand is a shaped type then the result "
-                                 "must be an extent tensor";
-  }
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // SizeToIndexOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])

mlir/test/Dialect/Shape/invalid.mlir

func @assuming_all_op_too_few_operands() {
   %w0 = shape.assuming_all
   return
 }
 
 // -----
 
 func @shape_of(%value_arg : !shape.value_shape,
                %shaped_arg : tensor<?x3x4xf32>) {
-  // expected-error@+1 {{if operand is of type `value_shape` then the result must be of type `shape` to propagate potential error shapes}}
+  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `shape` to propagate them}}
   %0 = shape.shape_of %value_arg : !shape.value_shape -> tensor<?xindex>
   return
 }
 
 // -----
 
-func @shape_of(%value_arg : !shape.value_shape,
-               %shaped_arg : tensor<?x3x4xf32>) {
-  // expected-error@+1 {{if operand is a shaped type then the result must be an extent tensor}}
-  %1 = shape.shape_of %shaped_arg : tensor<?x3x4xf32> -> !shape.shape
-  return
-}
-
-// -----
-
 func @rank(%arg : !shape.shape) {
-  // expected-error@+1 {{if operand is of type `shape` then the result must be of type `size` to propagate potential errors}}
+  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
   %0 = shape.rank %arg : !shape.shape -> index
   return
 }
 
 // -----
 
-func @get_extent_error_free(%arg : tensor<?xindex>) -> !shape.size {
-  %c0 = constant 0 : index
-  // expected-error@+1 {{if none of the operands can hold error values then the result must be of type `index`}}
-  %result = shape.get_extent %arg, %c0 : tensor<?xindex>, index -> !shape.size
-  return %result : !shape.size
-}
-
-// -----
-
 func @get_extent_error_possible(%arg : tensor<?xindex>) -> index {
   %c0 = shape.const_size 0
   // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> index
   return %result : index
 }
 
 // -----
 
-func @mul_error_free(%arg : index) -> !shape.size {
-  // expected-error@+1 {{if none of the operands can hold error values then the result must be of type `index`}}
-  %result = shape.mul %arg, %arg : index, index -> !shape.size
-  return %result : !shape.size
-}
-
-// -----
-
 func @mul_error_possible(%lhs : !shape.size, %rhs : index) -> index {
   // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
   %result = shape.mul %lhs, %rhs : !shape.size, index -> index
   return %result : index
 }