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

[KERNEL_GEN] Rename enable_unranked_kernels to enable_experimental_kernels.

Browse Source 
PiperOrigin-RevId: 348001576
Change-Id: If88f9062db8e30b33291cf9636dbc771310b47ad
This commit is contained in:
Alexander Belyaev 2020-12-17 05:07:48 -08:00 committed by TensorFlower Gardener
parent 800a3b1fb4
commit 4927f25808
40 changed files with 50 additions and 50 deletions

4
tensorflow/core/kernels/BUILD
View File

@ -16,8 +16,8 @@ load(
)
load(
"//tensorflow/core/kernels/mlir_generated:build_defs.bzl",
"if_mlir_experimental_kernels_enabled",
"if_mlir_generated_gpu_kernels_enabled",
"if_mlir_unranked_kernels_enabled",
)
# buildifier: disable=same-origin-load
@ -3347,7 +3347,7 @@ tf_kernel_library(
tf_kernel_library(
name = "cwise_op",
copts = if_mlir_generated_gpu_kernels_enabled(if_true = ["-DMLIR_GENERATED_GPU_KERNELS_ENABLED=1"]) +
if_mlir_unranked_kernels_enabled(if_true = ["-DMLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED=1"]),
if_mlir_experimental_kernels_enabled(if_true = ["-DMLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED=1"]),
prefix = "cwise_op",
deps = MATH_DEPS + if_mlir_generated_gpu_kernels_enabled(if_true = ["//tensorflow/core/kernels/mlir_generated:cwise_op"]),
)

2
tensorflow/core/kernels/cwise_op_add_1.cc
View File

@ -25,7 +25,7 @@ REGISTER6(BinaryOp, CPU, "AddV2", functor::add, float, Eigen::half, double,
REGISTER3(BinaryOp, GPU, "Add", functor::add, float, Eigen::half, double);
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(BinaryOp, GPU, "AddV2", functor::add, float, Eigen::half, double);
#endif

2
tensorflow/core/kernels/cwise_op_add_2.cc
View File

@ -33,7 +33,7 @@ REGISTER4(BinaryOp, GPU, "Add", functor::add, uint8, int64, complex64,
complex128);
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER5(BinaryOp, GPU, "AddV2", functor::add, uint8, uint32, int64, complex64,
complex128);
#else

4
tensorflow/core/kernels/cwise_op_bitwise_and.cc
View File

@ -22,7 +22,7 @@ REGISTER8(BinaryOp, CPU, "BitwiseAnd", functor::bitwise_and, int8, int16, int32,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER8(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, int8, int16, int32,
int64, uint8, uint16, uint32, uint64);
#else
@ -30,7 +30,7 @@ REGISTER8(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, int8, int16, int32,
REGISTER4(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, uint8, uint16,
uint32, uint64);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED
// !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow

4
tensorflow/core/kernels/cwise_op_bitwise_or.cc
View File

@ -22,7 +22,7 @@ REGISTER8(BinaryOp, CPU, "BitwiseOr", functor::bitwise_or, int8, int16, int32,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER8(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, int8, int16, int32,
int64, uint8, uint16, uint32, uint64);
#else
@ -30,7 +30,7 @@ REGISTER8(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, int8, int16, int32,
REGISTER4(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, uint8, uint16,
uint32, uint64);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED
// !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow

4
tensorflow/core/kernels/cwise_op_bitwise_xor.cc
View File

@ -22,7 +22,7 @@ REGISTER8(BinaryOp, CPU, "BitwiseXor", functor::bitwise_xor, int8, int16, int32,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER8(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, int8, int16, int32,
int64, uint8, uint16, uint32, uint64);
#else
@ -30,7 +30,7 @@ REGISTER8(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, int8, int16, int32,
REGISTER4(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, uint8, uint16,
uint32, uint64);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED
// !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow

2
tensorflow/core/kernels/cwise_op_ceil.cc
View File

@ -21,7 +21,7 @@ REGISTER4(UnaryOp, CPU, "Ceil", functor::ceil, float, Eigen::half, bfloat16,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Ceil", functor::ceil, float, Eigen::half, double);
#endif
#endif

2
tensorflow/core/kernels/cwise_op_conj.cc
View File

@ -27,7 +27,7 @@ REGISTER_KERNEL_BUILDER(
Name("Conj").Device(DEVICE_GPU).TypeConstraint<Variant>("T"),
UnaryVariantOp<GPUDevice, CONJ_VARIANT_UNARY_OP>);
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_KERNEL_BUILDER(
Name("Conj").Device(DEVICE_GPU).TypeConstraint<complex64>("T"),
UnaryOp<GPUDevice, functor::conj<complex64>>);

2
tensorflow/core/kernels/cwise_op_cos.cc
View File

@ -21,7 +21,7 @@ REGISTER6(UnaryOp, CPU, "Cos", functor::cos, float, Eigen::half, bfloat16,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Cos", functor::cos, float, Eigen::half, double);
#endif
#endif

2
tensorflow/core/kernels/cwise_op_equal_to_1.cc
View File

@ -28,7 +28,7 @@ REGISTER_KERNEL_BUILDER(
ApproximateEqualOp<CPUDevice, double>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER4(BinaryOp, GPU, "Equal", functor::equal_to, float, Eigen::half, double,
uint8);
#else

2
tensorflow/core/kernels/cwise_op_equal_to_2.cc
View File

@ -26,7 +26,7 @@ REGISTER6(BinaryOp, CPU, "Equal", functor::equal_to, int32, int64, complex64,
complex128, tstring, bool);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER6(BinaryOp, GPU, "Equal", functor::equal_to, int8, int16, int64,
complex64, complex128, bool);
#else

2
tensorflow/core/kernels/cwise_op_exp.cc
View File

@ -21,7 +21,7 @@ REGISTER6(UnaryOp, CPU, "Exp", functor::exp, float, Eigen::half, bfloat16,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER5(UnaryOp, GPU, "Exp", functor::exp, float, Eigen::half, double,
complex64, complex128);
#endif

2
tensorflow/core/kernels/cwise_op_floor.cc
View File

@ -21,7 +21,7 @@ REGISTER4(UnaryOp, CPU, "Floor", functor::floor, float, Eigen::half, bfloat16,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Floor", functor::floor, float, Eigen::half, double);
#endif
#endif

2
tensorflow/core/kernels/cwise_op_gpu_left_shift.cu.cc
View File

@ -20,7 +20,7 @@ limitations under the License.
namespace tensorflow {
namespace functor {
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
DEFINE_BINARY8(left_shift, int8, int16, int32, int64, uint8, uint16, uint32,
uint64);
#else

2
tensorflow/core/kernels/cwise_op_gpu_right_shift.cu.cc
View File

@ -20,7 +20,7 @@ limitations under the License.
namespace tensorflow {
namespace functor {
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
DEFINE_BINARY8(right_shift, int8, int16, int32, int64, uint8, uint16, uint32,
uint64);
#else

2
tensorflow/core/kernels/cwise_op_greater.cc
View File

@ -20,7 +20,7 @@ REGISTER9(BinaryOp, CPU, "Greater", functor::greater, float, Eigen::half,
double, int32, int64, uint8, int8, int16, bfloat16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER7(BinaryOp, GPU, "Greater", functor::greater, float, Eigen::half,
double, int64, uint8, int8, int16);
#else

2
tensorflow/core/kernels/cwise_op_greater_equal.cc
View File

@ -20,7 +20,7 @@ REGISTER9(BinaryOp, CPU, "GreaterEqual", functor::greater_equal, float,
Eigen::half, double, int32, int64, uint8, int8, int16, bfloat16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER7(BinaryOp, GPU, "GreaterEqual", functor::greater_equal, float,
Eigen::half, double, int64, uint8, int8, int16);
#else

2
tensorflow/core/kernels/cwise_op_imag.cc
View File

@ -28,7 +28,7 @@ REGISTER_COMPLEX(CPU, double, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_COMPLEX(GPU, float, complex64);
REGISTER_COMPLEX(GPU, double, complex128);
#endif

2
tensorflow/core/kernels/cwise_op_isinf.cc
View File

@ -21,7 +21,7 @@ REGISTER4(UnaryOp, CPU, "IsInf", functor::isinf, float, Eigen::half, bfloat16,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "IsInf", functor::isinf, float, Eigen::half, double);
#endif
#endif

2
tensorflow/core/kernels/cwise_op_left_shift.cc
View File

@ -22,7 +22,7 @@ REGISTER8(BinaryOp, CPU, "LeftShift", functor::left_shift, int8, int16, int32,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER8(BinaryOp, GPU, "LeftShift", functor::left_shift, int8, int16, int32,
int64, uint8, uint16, uint32, uint64);
#else

2
tensorflow/core/kernels/cwise_op_less.cc
View File

@ -22,7 +22,7 @@ REGISTER4(BinaryOp, CPU, "Less", functor::less, int64, uint8, int8, int16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER7(BinaryOp, GPU, "Less", functor::less, float, Eigen::half, double,
int64, uint8, int8, int16);
#else

2
tensorflow/core/kernels/cwise_op_less_equal.cc
View File

@ -23,7 +23,7 @@ REGISTER4(BinaryOp, CPU, "LessEqual", functor::less_equal, int64, uint8, int8,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER7(BinaryOp, GPU, "LessEqual", functor::less_equal, float, Eigen::half,
double, int64, uint8, int8, int16);
#else

2
tensorflow/core/kernels/cwise_op_log.cc
View File

@ -21,7 +21,7 @@ REGISTER6(UnaryOp, CPU, "Log", functor::log, float, Eigen::half, double,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Log", functor::log, float, Eigen::half, double);
#endif
#endif

4
tensorflow/core/kernels/cwise_op_logical_and.cc
View File

@ -20,10 +20,10 @@ REGISTER_KERNEL_BUILDER(Name("LogicalAnd").Device(DEVICE_CPU),
BinaryOp<CPUDevice, functor::logical_and>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_KERNEL_BUILDER(Name("LogicalAnd").Device(DEVICE_GPU),
BinaryOp<GPUDevice, functor::logical_and>);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED
// !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif
} // namespace tensorflow

2
tensorflow/core/kernels/cwise_op_logical_not.cc
View File

@ -20,7 +20,7 @@ REGISTER_KERNEL_BUILDER(Name("LogicalNot").Device(DEVICE_CPU),
UnaryOp<CPUDevice, functor::logical_not>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_KERNEL_BUILDER(Name("LogicalNot").Device(DEVICE_GPU),
UnaryOp<GPUDevice, functor::logical_not>);
#endif

4
tensorflow/core/kernels/cwise_op_logical_or.cc
View File

@ -20,10 +20,10 @@ REGISTER_KERNEL_BUILDER(Name("LogicalOr").Device(DEVICE_CPU),
BinaryOp<CPUDevice, functor::logical_or>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_KERNEL_BUILDER(Name("LogicalOr").Device(DEVICE_GPU),
BinaryOp<GPUDevice, functor::logical_or>);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED
// !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif
} // namespace tensorflow

2
tensorflow/core/kernels/cwise_op_mul_1.cc
View File

@ -31,7 +31,7 @@ REGISTER(BinaryOp, CPU, "Mul", functor::mul, int32);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER4(BinaryOp, GPU, "Mul", functor::mul, Eigen::half, float, double,
uint8);
#else

2
tensorflow/core/kernels/cwise_op_mul_2.cc
View File

@ -26,7 +26,7 @@ REGISTER6(BinaryOp, CPU, "Mul", functor::mul, int8, uint16, int16, int64,
complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER6(BinaryOp, GPU, "Mul", functor::mul, int8, uint16, int16, int64,
complex64, complex128);
#else

2
tensorflow/core/kernels/cwise_op_neg_1.cc
View File

@ -21,7 +21,7 @@ REGISTER4(UnaryOp, CPU, "Neg", functor::neg, int8, int16, int32, int64);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Neg", functor::neg, int8, int16, int64);
#endif

2
tensorflow/core/kernels/cwise_op_neg_2.cc
View File

@ -21,7 +21,7 @@ REGISTER6(UnaryOp, CPU, "Neg", functor::neg, Eigen::half, float, double,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER6(UnaryOp, GPU, "Neg", functor::neg, Eigen::half, float, double,
bfloat16, complex64, complex128);
#else

2
tensorflow/core/kernels/cwise_op_not_equal_to_1.cc
View File

@ -22,7 +22,7 @@ REGISTER7(BinaryOp, CPU, "NotEqual", functor::not_equal_to, uint16, uint32,
uint64, qint8, qint16, quint8, quint16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER4(BinaryOp, GPU, "NotEqual", functor::not_equal_to, float, Eigen::half,
double, uint8);
#else

2
tensorflow/core/kernels/cwise_op_not_equal_to_2.cc
View File

@ -26,7 +26,7 @@ REGISTER6(BinaryOp, CPU, "NotEqual", functor::not_equal_to, int32, int64,
complex64, complex128, tstring, bool);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER6(BinaryOp, GPU, "NotEqual", functor::not_equal_to, int8, int16, int64,
complex64, complex128, bool);
#else

2
tensorflow/core/kernels/cwise_op_real.cc
View File

@ -29,7 +29,7 @@ REGISTER_COMPLEX(CPU, double, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_COMPLEX(GPU, float, complex64);
REGISTER_COMPLEX(GPU, double, complex128);
#endif

2
tensorflow/core/kernels/cwise_op_right_shift.cc
View File

@ -22,7 +22,7 @@ REGISTER8(BinaryOp, CPU, "RightShift", functor::right_shift, int8, int16, int32,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER8(BinaryOp, GPU, "RightShift", functor::right_shift, int8, int16, int32,
int64, uint8, uint16, uint32, uint64);
#else

2
tensorflow/core/kernels/cwise_op_rsqrt.cc
View File

@ -21,7 +21,7 @@ REGISTER5(UnaryOp, CPU, "Rsqrt", functor::rsqrt, float, Eigen::half, double,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Rsqrt", functor::rsqrt, float, Eigen::half, double);
#endif
#endif

2
tensorflow/core/kernels/cwise_op_sign.cc
View File

@ -20,7 +20,7 @@ REGISTER8(UnaryOp, CPU, "Sign", functor::sign, float, double, int32, int64,
complex64, Eigen::half, bfloat16, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER6(UnaryOp, GPU, "Sign", functor::sign, float, Eigen::half, double,
int64, complex64, complex128);
#else

2
tensorflow/core/kernels/cwise_op_sin.cc
View File

@ -21,7 +21,7 @@ REGISTER6(UnaryOp, CPU, "Sin", functor::sin, float, Eigen::half, bfloat16,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Sin", functor::sin, float, Eigen::half, double);
#endif
#endif

2
tensorflow/core/kernels/cwise_op_sqrt.cc
View File

@ -21,7 +21,7 @@ REGISTER6(UnaryOp, CPU, "Sqrt", functor::sqrt, float, Eigen::half, double,
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED)
!defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER3(UnaryOp, GPU, "Sqrt", functor::sqrt, float, Eigen::half, double);
#endif
#endif

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

@ -3,8 +3,8 @@
load(
"//tensorflow/core/kernels/mlir_generated:build_defs.bzl",
"gen_kernel_library",
"if_mlir_experimental_kernels_enabled",
"if_mlir_generated_gpu_kernels_enabled",
"if_mlir_unranked_kernels_enabled",
)
load(
"//tensorflow:tensorflow.bzl",
@ -31,7 +31,7 @@ config_setting(
)
config_setting(
name = "mlir_use_unranked_kernels",
name = "mlir_experimental_kernels_enabled",
define_values = {"enable_unranked_kernels": "1"},
)
@ -70,7 +70,7 @@ filegroup(
name = "unary_kernel_srcs",
srcs = [
":enabled_unary_kernel_srcs",
] + if_mlir_unranked_kernels_enabled(
] + if_mlir_experimental_kernels_enabled(
if_true = [":experimental_unary_kernel_srcs"],
),
compatible_with = get_compatible_with_cloud(),
@ -182,7 +182,7 @@ tf_kernel_library(
# but we want to avoid building them if they are not needed.
deps = if_cuda_or_rocm([
":cwise_unary_op",
]) + if_mlir_unranked_kernels_enabled(
]) + if_mlir_experimental_kernels_enabled(
[
":cwise_binary_op",
],

4
tensorflow/core/kernels/mlir_generated/build_defs.bzl
View File

@ -77,9 +77,9 @@ def _gen_mlir_op(name, type):
# Kernels build rules.
################################################################################
def if_mlir_unranked_kernels_enabled(if_true, if_false = []):
def if_mlir_experimental_kernels_enabled(if_true, if_false = []):
return select({
"//tensorflow/core/kernels/mlir_generated:mlir_use_unranked_kernels": if_true,
"//tensorflow/core/kernels/mlir_generated:mlir_experimental_kernels_enabled": if_true,
"//conditions:default": if_false,
})