[shape] Add min and max ops

[jpienaar]

Summary

These are element-wise operations that operates on shapes with equal ranks.
Also add missing printer/parser for join operator.

[silvas]

Not sure if you want to add folders in this patch.

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

407

document the rank mismatch behavior explicitly (UB or propagated error?).

mlir/test/Dialect/Shape/ops.mlir

294

!shape.size test case?

[frgossen]

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

407

Also of size type (Shape_ShapeOrSizeType). Either documentation or the type should be updated.

424

minimum

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

let description = [{
     shape.join([], [?, ?]) -> [invalid]
     shape.join([1, ?], [2, ?, ?]) -> [invalid]
     ```
 
     `shape.join` also allows specifying an optional error string, that may be
     used to return an error to the user upon mismatch of dimensions.
 
     ```mlir
-    %c = shape.join %a, %b, error="<reason>" : !shape.shape
+    %c = shape.join %a, %b, error="<reason>" : !shape.shape, !shape.shape -> !shape.shape
     ```
   }];
 
   let arguments = (ins Shape_ShapeOrSizeType:$arg0, Shape_ShapeOrSizeType:$arg1,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeOrSizeType:$result);
+
+  let assemblyFormat = [{
+    $arg0 `,` $arg1 (`,` `error` `=` $error^)? attr-dict `:`
+      type($arg0) `,` type($arg1) `->` type($result)
+  }];
+}
+
+def Shape_MaxOp : Shape_Op<"max", [Commutative, NoSideEffect]> {
+  let summary = "Elementwise maximum";
+  let description = [{
+    Computes the elementwise maximum of two shapes with equal ranks. If either
+    operand is an error, then an error will be propagated to the result. If the
+    input types mismatch or the ranks do not match, then the result is an
+    error.
+  }];
+
+  let arguments = (ins Shape_ShapeOrSizeType:$lhs, Shape_ShapeOrSizeType:$rhs);
+  let results = (outs Shape_ShapeOrSizeType:$result);
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
+  }];
+}
+
+def Shape_MinOp : Shape_Op<"min", [Commutative, NoSideEffect]> {
+  let summary = "Elementwise minimum";
+  let description = [{
+    Computes the elementwise maximum of two shapes with equal ranks. If either
+    operand is an error, then an error will be propagated to the result. If the
+    input types mismatch or the ranks do not match, then the result is an
+    error.
+  }];
+
+  let arguments = (ins Shape_ShapeOrSizeType:$lhs, Shape_ShapeOrSizeType:$rhs);
+  let results = (outs Shape_ShapeOrSizeType:$result);
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
+  }];
 }
 
 def Shape_MulOp : Shape_Op<"mul", [Commutative, NoSideEffect]> {
   let summary = "Multiplication of sizes and indices";
   let description = [{
     Multiplies 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`,

mlir/test/Dialect/Shape/ops.mlir

func @test_constraints() {
   shape.assuming %w_all -> !shape.shape {
     %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
 
 func @eq_on_extent_tensors(%lhs : tensor<?xindex>,
-                                      %rhs : tensor<?xindex>) {
   %w0 = shape.cstr_eq %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
   return
 }
 
 func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
                                       %rhs : tensor<?xindex>) {
   %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
   return
func @rank(%shape : !shape.shape) -> !shape.size {
   return %rank : !shape.size
 }
 
 func @rank_on_extent_tensor(%shape : tensor<?xindex>) -> index {
   %rank = shape.rank %shape : tensor<?xindex> -> index
   return %rank : index
 }
 
-
 func @shape_eq_on_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 {
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }
 
 func @shape_eq_on_tensors(%a : tensor<?xindex>, %b : tensor<?xindex>) -> i1 {
   %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
   return %result : i1

 }
 
 func @is_broadcastable_on_shapes(%a : !shape.shape,
                                  %b : !shape.shape) -> i1 {
   %result = shape.is_broadcastable %a, %b
       : !shape.shape, !shape.shape
   return %result : i1
 }
+
+func @shape_upper_bounded_by_constant(%a: !shape.shape) -> !shape.shape {
+  %0 = shape.const_shape [4, 57, 92] : !shape.shape
+  %1 = shape.max %a, %0 : !shape.shape, !shape.shape -> !shape.shape
+  %2 = shape.join %0, %1, error="exceeded element-wise upper bound" :
+    !shape.shape, !shape.shape -> !shape.shape
+  return %2 : !shape.shape
+}
+
+func @shape_lower_bounded_by_constant(%a: !shape.shape) -> !shape.shape {
+  %0 = shape.const_shape [4, 57, 92] : !shape.shape
+  %1 = shape.min %a, %0 : !shape.shape, !shape.shape -> !shape.shape
+  %2 = shape.join %0, %1, error="lower bound element-wise exceeded" :
+    !shape.shape, !shape.shape -> !shape.shape
+  return %2 : !shape.shape
+}
+
+func @size_upper_bounded_by_constant(%a: !shape.size) -> !shape.size {
+  %0 = shape.const_size 5
+  %1 = shape.max %a, %0 : !shape.size, !shape.size -> !shape.size
+  %2 = shape.join %0, %1, error="exceeded element-wise upper bound" :
+    !shape.size, !shape.size -> !shape.size
+  return %2 : !shape.size
+}
+
+func @size_lower_bounded_by_constant(%a: !shape.size) -> !shape.size {
+  %0 = shape.const_size 9
+  %1 = shape.min %a, %0 : !shape.size, !shape.size -> !shape.size
+  %2 = shape.join %0, %1, error="lower bound element-wise exceeded" :
+    !shape.size, !shape.size -> !shape.size
+  return %2 : !shape.size
+}