def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", [Commutative]> {
let summary = "Determines if 2 shapes can be successfully broadcasted";
let description = [{
- Given 2 input shapes, return a witness specifying if they are broadcastable.
- This broadcastable follows the same logic as what shape.broadcast documents.
+ Given two input shapes or extent tensors, return a witness specifying if
+ they are broadcastable. This broadcastable follows the same logic as what
+ shape.broadcast documents.
"cstr" operations represent runtime assertions.
```
}];
- let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
+ let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
+ Shape_ShapeOrExtentTensorType:$rhs);
let results = (outs Shape_WitnessType:$result);
- let assemblyFormat = "$lhs `,` $rhs attr-dict";
+ let assemblyFormat = "$lhs `,` $rhs `:` type($lhs) `,` type($rhs) attr-dict";
let hasCanonicalizer = 1;
let hasFolder = 1;
// CHECK-NEXT: return
%cs0 = shape.const_shape [3, 1]
%cs1 = shape.const_shape [1, 5]
- %0 = shape.cstr_broadcastable %cs0, %cs1
+ %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
"consume.witness"(%0) : (!shape.witness) -> ()
return
}
// CHECK-NEXT: return
%cs0 = shape.const_shape [1, 3]
%cs1 = shape.const_shape [1, 5]
- %0 = shape.cstr_broadcastable %cs0, %cs1
+ %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
"consume.witness"(%0) : (!shape.witness) -> ()
return
}
// CHECK-NEXT: consume.witness
// CHECK-NEXT: return
%cs0 = shape.const_shape [1,3]
- %0 = shape.cstr_broadcastable %arg0, %cs0
+ %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape
"consume.witness"(%0) : (!shape.witness) -> ()
return
}
// CHECK-NEXT: shape.const_witness true
// CHECK-NEXT: consume.witness
// CHECK-NEXT: return
- %0 = shape.cstr_broadcastable %arg0, %arg0
+ %0 = shape.cstr_broadcastable %arg0, %arg0 : !shape.shape, !shape.shape
+ "consume.witness"(%0) : (!shape.witness) -> ()
+ return
+}
+
+// -----
+
+// Broadcastable canonicalization also works on extent tensors.
+// CHECK-LABEL: func @broadcastable_on_extent_tensors
+func @broadcastable_on_extent_tensors(%arg : tensor<?xindex>) {
+ // CHECK-NEXT: shape.const_witness true
+ // CHECK-NEXT: consume.witness
+ // CHECK-NEXT: return
+ %0 = shape.cstr_broadcastable %arg, %arg : tensor<?xindex>, tensor<?xindex>
"consume.witness"(%0) : (!shape.witness) -> ()
return
}
// CHECK-NEXT: return
%0 = shape.const_shape []
%1 = shape.shape_of %arg0 : tensor<?xf32>
- %2 = shape.cstr_broadcastable %0, %1
+ %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
"consume.witness"(%2) : (!shape.witness) -> ()
return
}
// CHECK-NEXT: return
%0 = shape.shape_of %arg0 : tensor<?xf32>
%1 = shape.shape_of %arg1 : tensor<?xf32>
- %2 = shape.cstr_broadcastable %0, %1
+ %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
"consume.witness"(%2) : (!shape.witness) -> ()
return
}
// CHECK-NEXT: return
%0 = shape.shape_of %arg1 : tensor<index>
%1 = shape.shape_of %arg0 : tensor<*xf32>
- %2 = shape.cstr_broadcastable %0, %1
+ %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
"consume.witness"(%2) : (!shape.witness) -> ()
return
}
func @test_constraints() {
%0 = shape.const_shape []
%1 = shape.const_shape [1, 2, 3]
- %w0 = shape.cstr_broadcastable %0, %1
+ %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
%w1 = shape.cstr_eq %0, %1
%w2 = shape.const_witness true
%w3 = shape.const_witness false
return
}
+func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
+ %rhs : tensor<?xindex>) {
+ %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
+ return
+}
+
func @test_mul(%lhs: !shape.size, %rhs: !shape.size) -> !shape.size {
%product = shape.mul %lhs, %rhs
return %product: !shape.size
// REPLACE: shape.assuming %[[WITNESS]]
// CANON-NEXT: test.source
// CANON-NEXT: return
- %0 = shape.cstr_broadcastable %arg0, %arg1
+ %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
%1 = shape.assuming %0 -> index {
%2 = "test.source"() : () -> (index)
shape.assuming_yield %2 : index
func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
// CANON-NEXT: test.source
// CANON-NEXT: return
- %0 = shape.cstr_broadcastable %arg0, %arg1
+ %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
%1 = shape.cstr_eq %arg0, %arg1
%2 = shape.assuming_all %0, %1
%3 = shape.assuming %0 -> index {