Enable the MLIR generated MulNoNan GPU kernels by default.
PiperOrigin-RevId: 385520919
Change-Id: Ib32ae5085843877557fc0953bbb7330e8a0e7dc6
diff --git a/tensorflow/core/kernels/cwise_op_gpu_mul.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_mul.cu.cc
index a87fca4..3d2cfda 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_mul.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_mul.cu.cc
@@ -23,8 +23,7 @@
DEFINE_BINARY5(mul, Eigen::half, int8, int16, int32, int64);
#endif
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
- !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
DEFINE_BINARY5(mul_no_nan, Eigen::half, float, double, complex64, complex128);
#endif
diff --git a/tensorflow/core/kernels/cwise_op_mul_1.cc b/tensorflow/core/kernels/cwise_op_mul_1.cc
index 6419a67..168e982 100644
--- a/tensorflow/core/kernels/cwise_op_mul_1.cc
+++ b/tensorflow/core/kernels/cwise_op_mul_1.cc
@@ -56,8 +56,7 @@
BinaryOp<CPUDevice, functor::mul<int32>>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
- !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
REGISTER5(BinaryOp, GPU, "MulNoNan", functor::mul_no_nan, Eigen::half, float,
double, complex64, complex128);
#endif
diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index 33be761..8f3620c 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -249,6 +249,7 @@
"gpu_op_maximum.cc",
"gpu_op_minimum.cc",
"gpu_op_mul.cc",
+ "gpu_op_mul_no_nan.cc",
"gpu_op_not_equal.cc",
"gpu_op_polygamma.cc",
"gpu_op_pow.cc",
@@ -259,7 +260,6 @@
"gpu_op_xdivy.cc",
"gpu_op_zeta.cc",
]) + if_mlir_generated_experimental_kernels_enabled([
- "gpu_op_mul_no_nan.cc",
"gpu_op_xlogy.cc",
"gpu_op_xlog1py.cc",
]),
@@ -286,6 +286,7 @@
":gpu_maximum_kernels",
":gpu_minimum_kernels",
":gpu_mul_kernels",
+ ":gpu_mul_no_nan_kernels",
":gpu_not_equal_kernels",
":gpu_polygamma_kernels",
":gpu_pow_kernels",
@@ -297,7 +298,6 @@
":gpu_zeta_kernels",
"//third_party/eigen3",
]) + if_mlir_generated_experimental_kernels_enabled([
- ":gpu_mul_no_nan_kernels",
":gpu_xlogy_kernels",
":gpu_xlog1py_kernels",
]),
