b70e4efb75
Summary: Generalize broadcastable trait to variadic operands. Update the documentation that still talked about element type as part of broadcastable trait (that bug was already fixed). Also rename Broadcastable to ResultBroadcastableShape to be more explicit that the trait affects the result shape (it is possible for op to allow broadcastable operands but not have result shape that is broadcast compatible with operands). Doing some intermediate work to have getBroadcastedType take an optional elementType as input and use that if specified, instead of the common element type of type1 and type2 in this function. Differential Revision: https://reviews.llvm.org/D72559
154 lines
5.9 KiB
MLIR
154 lines
5.9 KiB
MLIR
// RUN: mlir-opt %s -split-input-file -verify-diagnostics
|
|
|
|
// Verify that ops with broadcastable trait verifies operand and result type
|
|
// combinations and emits an error for invalid combinations.
|
|
|
|
func @broadcast_scalar_scalar_scalar(tensor<i32>, tensor<i32>) -> tensor<i32> {
|
|
^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
|
|
return %0 : tensor<i32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @broadcast_tensor_scalar_tensor(tensor<4xi32>, tensor<i32>) -> tensor<4xi32> {
|
|
^bb0(%arg0: tensor<4xi32>, %arg1: tensor<i32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4xi32>, tensor<i32>) -> tensor<4xi32>
|
|
return %0 : tensor<4xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check only one dimension has size 1
|
|
func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<3x1xi32>) -> tensor<4x3x2xi32> {
|
|
^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<3x1xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<3x1xi32>) -> tensor<4x3x2xi32>
|
|
return %0 : tensor<4x3x2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check multiple dimensions have size 1
|
|
func @broadcast_tensor_tensor_tensor(tensor<8x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x5xi32> {
|
|
^bb0(%arg0: tensor<8x1x6x1xi32>, %arg1: tensor<7x1x5xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<8x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x5xi32>
|
|
return %0 : tensor<8x7x6x5xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check leading unknown dimension
|
|
func @broadcast_tensor_tensor_tensor(tensor<?x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<?x7x6x5xi32> {
|
|
^bb0(%arg0: tensor<?x1x6x1xi32>, %arg1: tensor<7x1x5xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<?x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<?x7x6x5xi32>
|
|
return %0 : tensor<?x7x6x5xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check unknown dimension in the middle
|
|
func @broadcast_tensor_tensor_tensor(tensor<8x1x?x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x?x5xi32> {
|
|
^bb0(%arg0: tensor<8x1x?x1xi32>, %arg1: tensor<7x1x5xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<8x1x?x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x?x5xi32>
|
|
return %0 : tensor<8x7x?x5xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check incompatible vector and tensor result type
|
|
func @broadcast_scalar_vector_vector(tensor<4xf32>, tensor<4xf32>) -> vector<4xf32> {
|
|
^bb0(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>):
|
|
// expected-error @+1 {{cannot broadcast vector with tensor}}
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> vector<4xf32>
|
|
return %0 : vector<4xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check incompatible operand types with known dimension
|
|
func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<3x3xi32>) -> tensor<4x3x2xi32> {
|
|
^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<3x3xi32>):
|
|
// expected-error @+1 {{operands don't have broadcast-compatible shapes}}
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<3x3xi32>) -> tensor<4x3x2xi32>
|
|
return %0 : tensor<4x3x2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check incompatible result type with known dimension
|
|
func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<3x1xi32>) -> tensor<4x3x3xi32> {
|
|
^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<3x1xi32>):
|
|
// expected-error @+1 {{op result type '4x3x3' not broadcast compatible with broadcasted operands's shapes '4x3x2'}}
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<3x1xi32>) -> tensor<4x3x3xi32>
|
|
return %0 : tensor<4x3x3xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Check incompatible result type with known dimension
|
|
func @broadcast_tensor_tensor_tensor(tensor<8x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x1xi32> {
|
|
^bb0(%arg0: tensor<8x1x6x1xi32>, %arg1: tensor<7x1x5xi32>):
|
|
// expected-error @+1 {{op result type '8x7x6x1' not broadcast compatible with broadcasted operands's shapes '8x7x6x5'}}
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<8x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x1xi32>
|
|
return %0 : tensor<8x7x6x1xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @broadcast_tensor_tensor_tensor(tensor<2xi32>, tensor<2xi32>) -> tensor<*xi32> {
|
|
^bb0(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<*xi32>
|
|
return %0 : tensor<*xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @broadcast_tensor_tensor_tensor(tensor<4x3x2xi32>, tensor<?xi32>) -> tensor<4x3x2xi32> {
|
|
^bb0(%arg0: tensor<4x3x2xi32>, %arg1: tensor<?xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4x3x2xi32>, tensor<?xi32>) -> tensor<4x3x2xi32>
|
|
return %0 : tensor<4x3x2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Unranked operands but ranked result
|
|
func @broadcast_tensor_tensor_tensor(tensor<*xi32>, tensor<*xi32>) -> tensor<2xi32> {
|
|
^bb0(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<*xi32>, tensor<*xi32>) -> tensor<2xi32>
|
|
return %0 : tensor<2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Unranked operand and compatible ranked result
|
|
func @broadcast_tensor_tensor_tensor(tensor<3x2xi32>, tensor<*xi32>) -> tensor<4x3x2xi32> {
|
|
^bb0(%arg0: tensor<3x2xi32>, %arg1: tensor<*xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg0, %arg1) : (tensor<3x2xi32>, tensor<3x2xi32>, tensor<*xi32>) -> tensor<4x3x2xi32>
|
|
return %0 : tensor<4x3x2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @broadcast_tensor_tensor_tensor(tensor<3x2xi32>, tensor<*xi32>) -> tensor<2xi32> {
|
|
^bb0(%arg0: tensor<3x2xi32>, %arg1: tensor<*xi32>):
|
|
// expected-error @+1 {{op result type '2' not broadcast compatible with broadcasted operands's shapes '3x2'}}
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<3x2xi32>, tensor<*xi32>) -> tensor<2xi32>
|
|
return %0 : tensor<2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @broadcast_tensor_tensor_tensor(tensor<?x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x5xi32> {
|
|
^bb0(%arg0: tensor<?x1x6x1xi32>, %arg1: tensor<7x1x5xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<?x1x6x1xi32>, tensor<7x1x5xi32>) -> tensor<8x7x6x5xi32>
|
|
return %0 : tensor<8x7x6x5xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @broadcastDifferentResultType(tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> {
|
|
^bb0(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>):
|
|
%0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
|
|
return %0 : tensor<4xi1>
|
|
}
|