[mlir] Add shape.is_broadcastable.

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.

Author: tpopp

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

@@ -181,6 +181,27 @@ def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
   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 = [{

mlir/test/Dialect/Shape/ops.mlir

@@ -260,3 +260,17 @@ func @any_on_extent_tensors(%a : tensor<?xindex>,
       : 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
+}