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

Register MLIR generated kernel for Complex.

Browse Source 
PiperOrigin-RevId: 351954046
Change-Id: I9df849dc53fae72fbdb6dc56d82c0b36a2ded230
This commit is contained in:
Adrian Kuegel 2021-01-15 00:03:12 -08:00 committed by TensorFlower Gardener
parent 7180803d3a
commit 7c8538cbd7
6 changed files with 29 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/compiler/mlir/xla/tests/legalize-tf.mlir
View File

@ -928,7 +928,7 @@ func @broadcast_to(%arg0: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
// CHECK-LABEL: func @complex
func @complex(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xcomplex<f32>> {
// CHECK: "mhlo.complex"
// CHECK: chlo.broadcast_complex
%1 = "tf.Complex"(%arg0, %arg1) : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xcomplex<f32>>
return %1 : tensor<3xcomplex<f32>>
}

3
tensorflow/compiler/mlir/xla/transforms/legalize_tf_patterns.td
View File

@ -92,6 +92,7 @@ class DirectBinaryPat<Op FromOp, Op ToOp>
foreach fromToBinPair = [[TF_AddV2Op, HLOClient_BroadcastAddOp],
[TF_Atan2Op, HLOClient_BroadcastAtan2Op],
[TF_ComplexOp, HLOClient_BroadcastComplexOp],
[TF_DivOp, HLOClient_BroadcastDivOp],
[TF_LeftShiftOp, HLOClient_BroadcastShiftLeftOp],
[TF_MaximumOp, HLOClient_BroadcastMaxOp],
@ -111,8 +112,6 @@ def LowerRightShiftSigned :
// TODO(hinsu): Lower unsigned types to HLO_ShiftRightLogical once the HLO op
// supports unsigned integers.
def : Pat<(TF_ComplexOp $r, $i), (HLO_ComplexOp $r, $i)>;
// Performs a substitution of FloorDiv, pseudo code below:
//
// return floor(div(x, y))

3
tensorflow/core/kernels/cwise_op_complex.cc
View File

@ -27,9 +27,12 @@ REGISTER_COMPLEX(CPU, float, complex64);
REGISTER_COMPLEX(CPU, double, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_COMPLEX(GPU, float, complex64);
REGISTER_COMPLEX(GPU, double, complex128);
#endif
#endif
#undef REGISTER_COMPLEX
} // namespace tensorflow

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

@ -50,6 +50,7 @@ filegroup(
"gpu_op_asin.cc",
"gpu_op_atan.cc",
"gpu_op_ceil.cc",
"gpu_op_complex.cc",
"gpu_op_conj.cc",
"gpu_op_cos.cc",
"gpu_op_exp.cc",
@ -112,6 +113,7 @@ tf_kernel_library(
":asin_kernels",
":atan_kernels",
":ceil_kernels",
":complex_kernels",
":conj_kernels",
":cos_kernels",
":exp_kernels",

19
tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc
View File

@ -469,6 +469,25 @@ GENERATE_DEFAULT_TESTS(BitwiseXor,
GENERATE_DEFAULT_TESTS(BitwiseXor,
/*test_name=*/Int64, int64, int64, baseline_bitwise_xor)
/// Test `tf.Complex`.
template <typename T>
std::complex<T> baseline_complex(T lhs, T rhs) {
return std::complex<T>(lhs, rhs);
}
GENERATE_DEFAULT_TESTS_2(
Complex,
/*test_name=*/C64, float, float, std::complex<float>, std::complex<float>,
test::DefaultInput<float>(), test::DefaultInput<float>(), baseline_complex,
test::GpuOpsTestConfig().ExpectStrictlyEqual().AddTout())
GENERATE_DEFAULT_TESTS_2(
Complex,
/*test_name=*/C128, double, double, std::complex<double>,
std::complex<double>, test::DefaultInput<double>(),
test::DefaultInput<double>(), baseline_complex,
test::GpuOpsTestConfig().ExpectStrictlyEqual().AddTout())
/// Test `tf.Div`.
template <typename T>

6
tensorflow/core/kernels/mlir_generated/gpu_op_complex.cc
View File

@ -20,10 +20,10 @@ limitations under the License.
namespace tensorflow {
GENERATE_UNARY_KERNEL2(Complex, f32, DT_COMPLEX64, std::complex<float>, float);
GENERATE_BINARY_KERNEL2(Complex, f32, DT_COMPLEX64, std::complex<float>, float);
REGISTER_COMPLEX_KERNEL(Complex, f32, std::complex<float>, float);
GENERATE_UNARY_KERNEL2(Complex, f64, DT_COMPLEX128, std::complex<double>,
double);
GENERATE_BINARY_KERNEL2(Complex, f64, DT_COMPLEX128, std::complex<double>,
double);
REGISTER_COMPLEX_KERNEL(Complex, f64, std::complex<double>, double);
} // namespace tensorflow