[MLIR] Forward extent tensors through shape.broadcast.

[tpopp]

Summary

Allow shape.broadcast to return extent tensors when all inputs are extent tensors.

[herhut]

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

86

DBC: Can you add to the comment and stress the fact that the op may only return a tensor if we statically know that it cannot fail, either because it is guarded by a corresponding cstr_broadcastable or the construction otherwise ensures this?

[frgossen]

Thanks

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

2. Else the input array with number of dimensions smaller than the max
 
            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]
 
     Op has an optional string attribute for the error case where there is no
     broadcastable output shape possible for the given inputs.
   }];
 
   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;
+
+  let verifier = [{ return ::verifyShapeOrExtentTensorOp(*this); }];
 }
 
 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
     as the `shape` attribute. The number of these values equals the shape's
     rank.

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
 }
 
 // -----
 
 // 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/invalid.mlir

 // -----
 
 func @num_elements_error_possible(%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 @broadcast_error_possible(%arg0 : !shape.shape, %arg1 : !shape.shape) -> tensor<?xindex> {
+  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `shape` to propagate them}}
+  %result = shape.broadcast %arg0, %arg1 : !shape.shape, !shape.shape -> tensor<?xindex>
+  return %result : tensor<?xindex>
+}
+
+
+// -----
+
+func @broadcast_error_possible(%arg0 : !shape.shape, %arg1 : tensor<?xindex>) -> tensor<?xindex> {
+  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `shape` to propagate them}}
+  %result = shape.broadcast %arg0, %arg1 : !shape.shape, tensor<?xindex> -> tensor<?xindex>
+  return %result : tensor<?xindex>
+}

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_broadcast_extents() -> tensor<?xindex> {
+  %0 = shape.const_shape [10, 1, 57, 92] : tensor<?xindex>
+  %1 = shape.const_shape [4, 57, 92] : tensor<?xindex>
+  %2 = shape.broadcast %0, %1 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+  return %2 : tensor<?xindex>
+}
+
 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
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }