Add missing lowering step for IsFiniteOp.
Also add a BUILD target for generating the GPU kernel. PiperOrigin-RevId: 334993362
This commit is contained in:
parent
5f303440da
commit
3eb767b43d
|
@ -506,6 +506,7 @@ void populateHLOToLHLOConversionPattern(
|
||||||
HloToLhloOpConverter<mhlo::GatherOp>,
|
HloToLhloOpConverter<mhlo::GatherOp>,
|
||||||
HloToLhloOpConverter<mhlo::ImagOp>,
|
HloToLhloOpConverter<mhlo::ImagOp>,
|
||||||
HloToLhloOpConverter<mhlo::IotaOp>,
|
HloToLhloOpConverter<mhlo::IotaOp>,
|
||||||
|
HloToLhloOpConverter<mhlo::IsFiniteOp>,
|
||||||
HloToLhloOpConverter<mhlo::LogOp>,
|
HloToLhloOpConverter<mhlo::LogOp>,
|
||||||
HloToLhloOpConverter<mhlo::MaxOp>,
|
HloToLhloOpConverter<mhlo::MaxOp>,
|
||||||
HloToLhloOpConverter<mhlo::MinOp>,
|
HloToLhloOpConverter<mhlo::MinOp>,
|
||||||
|
|
|
@ -595,9 +595,20 @@ func @custom_call(%arg0: memref<2x2xf32>, %arg1: memref<2x3xf32>, %result: memre
|
||||||
%arg0_tensor = tensor_load %arg0 : memref<2x2xf32>
|
%arg0_tensor = tensor_load %arg0 : memref<2x2xf32>
|
||||||
%arg1_tensor = tensor_load %arg1 : memref<2x3xf32>
|
%arg1_tensor = tensor_load %arg1 : memref<2x3xf32>
|
||||||
// BOTH: "lmhlo.custom_call"([[ARG0]], [[ARG1]], %{{.*}}) {backend_config = "", call_target_name = "foo", has_side_effect = false}
|
// BOTH: "lmhlo.custom_call"([[ARG0]], [[ARG1]], %{{.*}}) {backend_config = "", call_target_name = "foo", has_side_effect = false}
|
||||||
%result_tensor = "mhlo.custom_call"(%arg0_tensor, %arg1_tensor)
|
%result_tensor = "mhlo.custom_call"(%arg0_tensor, %arg1_tensor)
|
||||||
{backend_config = "", call_target_name = "foo", has_side_effect = false}
|
{backend_config = "", call_target_name = "foo", has_side_effect = false}
|
||||||
: (tensor<2x2xf32>, tensor<2x3xf32>) -> tensor<4x4xf16>
|
: (tensor<2x2xf32>, tensor<2x3xf32>) -> tensor<4x4xf16>
|
||||||
tensor_store %result_tensor, %result: memref<4x4xf16>
|
tensor_store %result_tensor, %result: memref<4x4xf16>
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ----
|
||||||
|
|
||||||
|
// BOTH-LABEL: func @isfinite
|
||||||
|
func @isfinite(%arg0: memref<2x2xf32>, %result: memref<2x2xi1>) {
|
||||||
|
%arg0_tensor = tensor_load %arg0 : memref<2x2xf32>
|
||||||
|
// BOTH: "lmhlo.is_finite"(%{{.*}}, %{{.*}})
|
||||||
|
%result_tensor = "mhlo.is_finite"(%arg0_tensor) : (tensor<2x2xf32>) -> tensor<2x2xi1>
|
||||||
|
tensor_store %result_tensor, %result: memref<2x2xi1>
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue