[mlir] Add shape.is_broadcastable.

[tpopp]

Summary

This op returns a boolean value indicating whether 2 ops are
broadcastable or not. This follows the same logic as the other ops with
broadcast in their names in the shape dialect.

Concretely, shape.is_broadcastable returning true implies that
shape.broadcast will not give an error, and shape.cstr_broadcastable
will not result in an assertion failure. Similarly, false implies an
error or assertion failure.

[jpienaar]

Thanks

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

188

Could you document error case here too? (You do in change summary)

202

attr-dict before : (we are inconsistent here, and I should fix that, but that closer correspondence to generic & std and reads more naturally IMHO)

[tpopp]

Somehow I submitted this already, but the commit detached from this review. I am abandoning this change, so it goes away, but I will handle feedback in a follow up commit.

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

def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
   }];
 
   let arguments = (ins IndexTensor:$input);
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "attr-dict $input `:` type($input)";
 }
 
+def Shape_IsBroadcastableOp : Shape_Op<"is_broadcastable", [Commutative]> {
+  let summary = "Determines if 2 shapes can be successfully broadcasted";
+  let description = [{
+    Given two input shapes or extent tensors, return a predicate specifying if
+    they are broadcastable. This broadcastable follows the same logic as what
+    shape.broadcast documents.
+
+    Example:
+    ```mlir
+    %true = shape.is_broadcastable [2,2], [3,1,2]
+    %false = shape.is_broadcastable [2,2], [3,2]
+    ```
+  }];
+
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
+                       Shape_ShapeOrExtentTensorType:$rhs);
+  let results = (outs I1:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs `:` type($lhs) `,` type($rhs) attr-dict";
+}
+
 def Shape_RankOp : Shape_Op<"rank", [NoSideEffect]> {
   let summary = "Gets the rank of a shape";
   let description = [{
     Returns the rank of the shape or extent tensor, i.e. the number of extents.
   }];
 
   let arguments = (ins Shape_ShapeOrExtentTensorType:$shape);
   let results = (outs Shape_SizeOrIndexType:$rank);

mlir/test/Dialect/Shape/ops.mlir

 
 func @any_on_extent_tensors(%a : tensor<?xindex>,
                             %b : tensor<?xindex>,
                             %c : tensor<?xindex>) -> tensor<?xindex> {
   %result = shape.any %a, %b, %c
       : tensor<?xindex>, tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
   return %result : tensor<?xindex>
 }
+
+func @is_broadcastable_on_extent_tensors(%a : tensor<?xindex>,
+                                         %b : tensor<?xindex>) -> i1 {
+  %result = shape.is_broadcastable %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
+}