[MLIR][KernelGen] Register all JIT kernels with the appropriate label
PiperOrigin-RevId: 407043105
Change-Id: I707ad9defaf3500b00d5bbfb0732b84f456ee4e2
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_abs.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_abs.cc
index 0d5bcaa..41bed39 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_abs.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_abs.cc
@@ -26,8 +26,8 @@
// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(Abs, DT_INT8);
-GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(Abs, DT_INT16);
+GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(Abs, DT_INT8);
+GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(Abs, DT_INT16);
#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_div.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_div.cc
index 5dd82fa..94db1de 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_div.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_div.cc
@@ -30,9 +30,9 @@
// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Div, DT_INT8);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Div, DT_UINT32);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Div, DT_UINT64);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Div, DT_INT8);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Div, DT_UINT32);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Div, DT_UINT64);
#endif
REGISTER_ALIASED_GPU_KERNEL(RealDiv, Div, DT_HALF, DT_HALF);
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_floor_div.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_floor_div.cc
index a7b7acf..e497906 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_floor_div.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_floor_div.cc
@@ -20,10 +20,11 @@
// The aliased kernels for data types DT_UINT8 and DT_UINT16 are registered in
// gpu_op_div.cc, because they alias the Div kernel.
+// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(FloorDiv, DT_INT8);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(FloorDiv, DT_UINT32);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(FloorDiv, DT_UINT64);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(FloorDiv, DT_INT8);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(FloorDiv, DT_UINT32);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(FloorDiv, DT_UINT64);
#endif
GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(FloorDiv, DT_INT16);
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_maximum.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_maximum.cc
index 93ca6c9..a2cbca4 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_maximum.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_maximum.cc
@@ -28,10 +28,10 @@
// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Maximum, DT_INT8);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Maximum, DT_UINT16);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Maximum, DT_UINT32);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Maximum, DT_UINT64);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_INT8);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_UINT16);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_UINT32);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Maximum, DT_UINT64);
#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_minimum.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_minimum.cc
index ec9eff6..5567ca6 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_minimum.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_minimum.cc
@@ -28,10 +28,10 @@
// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Minimum, DT_INT8);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Minimum, DT_UINT16);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Minimum, DT_UINT32);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Minimum, DT_UINT64);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_INT8);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_UINT16);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_UINT32);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Minimum, DT_UINT64);
#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc
index f04a104..98c5ba6 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc
@@ -24,8 +24,8 @@
// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_INT8);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_INT16);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Pow, DT_INT8);
+GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Pow, DT_INT16);
#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_rint.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_rint.cc
index 49e0db2..b9d34e8 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_rint.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_rint.cc
@@ -18,11 +18,11 @@
namespace tensorflow {
-// This kernel is jitted. Leave as experimental until we want to test jitted
-// kernels in more cases.
+// This kernel is JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(Rint, DT_HALF);
#endif
+
GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(Rint, DT_FLOAT);
GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(Rint, DT_DOUBLE);
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_sign.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_sign.cc
index cea612e..e4ad9df 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_sign.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_sign.cc
@@ -28,8 +28,8 @@
// These kernels are JIT-compiled.
#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
-GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(Sign, DT_INT8);
-GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(Sign, DT_INT16);
+GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(Sign, DT_INT8);
+GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(Sign, DT_INT16);
#endif
} // namespace tensorflow
