experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add missing lowering step for IsFiniteOp.

Browse Source 
Also add a BUILD target for generating the GPU kernel.

PiperOrigin-RevId: 334993362
Change-Id: I712f9d138c7050503d98a53412cd9c923a1ff60d
This commit is contained in:
Adrian Kuegel 2020-10-02 03:07:56 -07:00 committed by TensorFlower Gardener
parent 81d9fcbf83
commit fed53e76a7
3 changed files with 25 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
View File

@ -506,6 +506,7 @@ void populateHLOToLHLOConversionPattern(
HloToLhloOpConverter<mhlo::GatherOp>,
HloToLhloOpConverter<mhlo::ImagOp>,
HloToLhloOpConverter<mhlo::IotaOp>,
HloToLhloOpConverter<mhlo::IsFiniteOp>,
HloToLhloOpConverter<mhlo::LogOp>,
HloToLhloOpConverter<mhlo::MaxOp>,
HloToLhloOpConverter<mhlo::MinOp>,

13
tensorflow/compiler/mlir/hlo/tests/hlo-legalize-to-lhlo.mlir
View File

@ -595,9 +595,20 @@ func @custom_call(%arg0: memref<2x2xf32>, %arg1: memref<2x3xf32>, %result: memre
%arg0_tensor = tensor_load %arg0 : memref<2x2xf32>
%arg1_tensor = tensor_load %arg1 : memref<2x3xf32>
// 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}
: (tensor<2x2xf32>, tensor<2x3xf32>) -> tensor<4x4xf16>
tensor_store %result_tensor, %result: memref<4x4xf16>
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
}

12
tensorflow/core/kernels/mlir_generated/BUILD
View File

@ -243,6 +243,18 @@ gen_kernel_library(
unroll_factors = "4",
)
gen_kernel_library(
name = "isfinite",
same_shape = "0,1",
tile_size = "256",
types = [
"f16",
"f32",
"f64",
],
unroll_factors = "4",
)
gen_kernel_library(
name = "log",
same_shape = "0,1",