[MLIR][Shape] Allow `shape.add` to operate on indices

[frgossen]

Summary

Depends On D84439

[jpienaar]

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

69

Why is attr-dict at the end? Mostly we have it before the ':' in std etc (else we have add %x, %y : index, index - > index {foo = "blah" })

mlir/test/Dialect/Shape/invalid.mlir

133

I removed this case from verification for now as it makes progressive lowering require cast ops still. We can discuss this more tomorrow.

[frgossen]

Address comments

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

 //===----------------------------------------------------------------------===//
 // Shape op definitions
 //===----------------------------------------------------------------------===//
 
 // Base class for the operation in this dialect
 class Shape_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<ShapeDialect, mnemonic, traits>;
 
-def Shape_AddOp : Shape_Op<"add", [Commutative, SameOperandsAndResultType]> {
-  let summary = "Addition of sizes";
+def Shape_AddOp : Shape_Op<"add", [Commutative, NoSideEffect]> {
+  let summary = "Addition of sizes and indices";
   let description = [{
-    Adds two valid sizes as follows:
-    * lhs + rhs = unknown if either lhs or rhs unknown;
-    * lhs + rhs = (int)lhs + (int)rhs if known;
+    Adds two sizes or indices. If either operand is an error it will be
+    propagated to the result. The operands can be of type `size` or `index`. If
+    at least one of the operands can hold an error, i.e. if it is of type `size`,
+    then also the result must be of type `size`.
   }];
 
-  let arguments = (ins Shape_SizeType:$lhs, Shape_SizeType:$rhs);
-  let results = (outs Shape_SizeType:$result);
+  let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
+  let results = (outs Shape_SizeOrIndexType:$result);
 
-  let assemblyFormat = "$lhs `,` $rhs attr-dict";
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
+  }];
+
+  let verifier = [{ return verifySizeOrIndexOp(*this); }];
 }
 
 def Shape_BroadcastOp : Shape_Op<"broadcast", [Commutative]> {
   let summary = "Returns the broadcasted output shape of two inputs";
   let description = [{
     Computes the broadcasted output shape following:
     1. If any inputs are unranked, output is unranked;
     2. Else the input array with number of dimensions smaller than the max

mlir/test/Dialect/Shape/invalid.mlir

 func @rank(%arg : !shape.shape) {
   // 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_possible(%arg : tensor<?xindex>) -> index {
+func @get_extent(%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_possible(%lhs : !shape.size, %rhs : index) -> index {
+func @mul(%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
 }
 
 // -----
 
-func @num_elements_error_possible(%arg : !shape.shape) -> index {
+func @num_elements(%arg : !shape.shape) -> 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.num_elements %arg : !shape.shape -> index
   return %result : index
 }
 
+// -----
+
+func @add(%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.add %lhs, %rhs : !shape.size, index -> index
+  return %result : index
+}
+

mlir/test/Dialect/Shape/ops.mlir

func @mul(%size_arg : !shape.size, %index_arg : index) {
   %size_prod = shape.mul %size_arg, %size_arg
       : !shape.size, !shape.size -> !shape.size
   %index_prod = shape.mul %index_arg, %index_arg : index, index -> index
   %mixed_prod = shape.mul %size_arg, %index_arg
       : !shape.size, index -> !shape.size
   return
 }
 
+func @add(%size_arg : !shape.size, %index_arg : index) {
+  %size_sum = shape.add %size_arg, %size_arg
+      : !shape.size, !shape.size -> !shape.size
+  %index_sum = shape.add %index_arg, %index_arg : index, index -> index
+  %mixed_sum = shape.add %size_arg, %index_arg
+      : !shape.size, index -> !shape.size
+  return
+}
+
 func @const_size() {
   // CHECK: %c1 = shape.const_size 1
   // CHECK: %c2 = shape.const_size 2
   // CHECK: %c2_0 = shape.const_size 2
   %0 = shape.const_size 1
   %1 = shape.const_size 2
   %2 = shape.const_size 2
   return