[MLIR][Shape] Allow unsafe `shape.broadcast`

[frgossen]

Summary

In a context in which shape.broadcast is known not to produce an error value,
we want it to operate solely on extent tensors. The operation's behavior is
then undefined in the error case as the result type cannot hold this value.

[herhut]

mlir/test/Dialect/Shape/canonicalize.mlir

69

I would disallow this case. You can only have a tensor result of both inputs are tensors. Otherwise, there is no guarantee that the inputs will not be an error.

[frgossen]

mlir/test/Dialect/Shape/canonicalize.mlir

69

If I understand the semantics correctly, even if both inputs are of type tensor<?xindex> the result can be an error: [1,2], [3,4] -> error
That is ofc looking locally at this operation only.
We can disallow this case but we will still have undefined behaviour in some cases.

[herhut]

mlir/test/Dialect/Shape/canonicalize.mlir

69

Yes, you can only have a tensor result if both inputs are tensors but having tensor inputs is not sufficient to mandate a tensor output. Still, we should disallow cases that are statically wrong.

[frgossen]

Address comment

[herhut]

Thanks. With one more test this is good to land!

mlir/test/Dialect/Shape/canonicalize.mlir

65

Could you add the tensor, tensor -> shape case as well? Just to make sure it stays supported.

[frgossen]

Rebase, add test

[herhut]

Thank you!

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

def Shape_AddOp : Shape_Op<"add", [Commutative, SameOperandsAndResultType]> {
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
 }
 
 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
-       input dimension, has 1’s prepended to its shapes and the output shape is
-       calculated as follows:
+    Returns the broadcasted shape for two input shapes or extent tensors. Both
+    operands and the result type can be of type `shape.shape` or
+    `tensor<?xindex>`. If the two operand shapes are of different rank the
+    smaller one is padded with 1's from the left. The resulting broadcasted
+    shape is then defined as
 
-           output[i] = lhs[i] if lhs[i] == rhs[i] or rhs[i] is unknown/undefined
-                     = rhs[i] if lhs[i] is unknown/undefined
-                     = lhs[i] if rhs[i] == 1
-                     = rhs[i] if lhs[i] == 1
-                     = error  if lhs[i] != rhs[i]
+        result[i] = lhs[i] if lhs[i] == rhs[i]
+                  = rhs[i] if lhs[i] == 1.
 
-    Op has an optional string attribute for the error case where there is no
-    broadcastable output shape possible for the given inputs.
+    In case the resulting shape is undefined, i.e. if corresponding extents are
+    different from each other but none is 1, the result is an error shape.
+    Likewise error values are propagated if any of the operands holds an error
+    value. If the result type is an extent tensor (and can therefore not hold
+    the error value) the behavior is undefined. The optional string attribute
+    can be used to describe the error case.
   }];
 
   let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
                        Shape_ShapeOrExtentTensorType:$rhs,
                        OptionalAttr<StrAttr>:$error);
-  let results = (outs Shape_ShapeType:$result);
+  let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
-  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs)";
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
+  }];
 
   let hasFolder = 1;
 }
 
 def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
   let summary = "Creates a constant shape or extent tensor";
   let description = [{
     Creates a constant shape or extent tensor. The individual extents are given

mlir/test/Dialect/Shape/canonicalize.mlir

 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.const_shape [7, 2] : !shape.shape
   %0 = shape.const_shape [1, 2] : !shape.shape
   %1 = shape.const_shape [7, 1] : !shape.shape
-  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
 
+// Basic case including extent tensors.
+// CHECK-LABEL: @broadcast
+func @broadcast() -> tensor<?xindex> {
+  // CHECK: shape.const_shape [7, 2] : tensor<?xindex>
+  %0 = shape.const_shape [1, 2] : !shape.shape
+  %1 = shape.const_shape [7, 1] : tensor<?xindex>
+  %2 = shape.broadcast %0, %1 : !shape.shape, tensor<?xindex> -> tensor<?xindex>
+  return %2 : tensor<?xindex>
+}
+
+// -----
+
 // Rhs is a scalar.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {
   // CHECK: return %arg0
   %0 = shape.const_shape [] : !shape.shape
-  %1 = shape.broadcast %arg0, %0 : !shape.shape, !shape.shape
+  %1 = shape.broadcast %arg0, %0 : !shape.shape, !shape.shape -> !shape.shape
   return %1 : !shape.shape
 }
 
 // -----
 
 // Lhs is a scalar.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {
   // CHECK: return %arg0
   %0 = shape.const_shape [] : !shape.shape
-  %1 = shape.broadcast %0, %arg0 : !shape.shape, !shape.shape
+  %1 = shape.broadcast %0, %arg0 : !shape.shape, !shape.shape -> !shape.shape
   return %1 : !shape.shape
 }
 
 // -----
 
 // Lhs is a scalar and rhs is constant.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: %[[CST:.*]] = shape.const_shape [1, 2, 3] : !shape.shape
   // CHECK: return %[[CST]]
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
-  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
 
 // Incompatible shapes. No folding.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.broadcast
   %0 = shape.const_shape [2] : !shape.shape
   %1 = shape.const_shape [7] : !shape.shape
-  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {

mlir/test/Dialect/Shape/ops.mlir

func @test_shape_num_elements_fixed() {
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @test_broadcast_fixed() {
   %0 = shape.const_shape [10, 1, 57, 92] : !shape.shape
   %1 = shape.const_shape [4, 57, 92] : !shape.shape
-  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_fixed() {
   %0 = shape.const_shape [4, 57, 92] : !shape.shape
   %1 = shape.const_shape [4, 57, 92] : !shape.shape
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape