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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/kernels/BUILD
View File

@ -16,8 +16,8 @@ load(
) )
load( load(
"//tensorflow/core/kernels/mlir_generated:build_defs.bzl", "//tensorflow/core/kernels/mlir_generated:build_defs.bzl",
"if_mlir_experimental_kernels_enabled",
"if_mlir_generated_gpu_kernels_enabled", "if_mlir_generated_gpu_kernels_enabled",
"if_mlir_unranked_kernels_enabled",
) )
# buildifier: disable=same-origin-load # buildifier: disable=same-origin-load
@ -3347,7 +3347,7 @@ tf_kernel_library(
tf_kernel_library( tf_kernel_library(
name = "cwise_op", name = "cwise_op",
copts = if_mlir_generated_gpu_kernels_enabled(if_true = ["-DMLIR_GENERATED_GPU_KERNELS_ENABLED=1"]) + 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", prefix = "cwise_op",
deps = MATH_DEPS + if_mlir_generated_gpu_kernels_enabled(if_true = ["//tensorflow/core/kernels/mlir_generated: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); REGISTER3(BinaryOp, GPU, "Add", functor::add, float, Eigen::half, double);
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(BinaryOp, GPU, "AddV2", functor::add, float, Eigen::half, double);
#endif #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); complex128);
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER5(BinaryOp, GPU, "AddV2", functor::add, uint8, uint32, int64, complex64,
complex128); complex128);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER8(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, int8, int16, int32,
int64, uint8, uint16, uint32, uint64); int64, uint8, uint16, uint32, uint64);
#else #else
@ -30,7 +30,7 @@ REGISTER8(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, int8, int16, int32,
REGISTER4(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, uint8, uint16, REGISTER4(BinaryOp, GPU, "BitwiseAnd", functor::bitwise_and, uint8, uint16,
uint32, uint64); uint32, uint64);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED || #endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED // !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow } // 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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER8(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, int8, int16, int32,
int64, uint8, uint16, uint32, uint64); int64, uint8, uint16, uint32, uint64);
#else #else
@ -30,7 +30,7 @@ REGISTER8(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, int8, int16, int32,
REGISTER4(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, uint8, uint16, REGISTER4(BinaryOp, GPU, "BitwiseOr", functor::bitwise_or, uint8, uint16,
uint32, uint64); uint32, uint64);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED || #endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED // !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow } // 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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER8(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, int8, int16, int32,
int64, uint8, uint16, uint32, uint64); int64, uint8, uint16, uint32, uint64);
#else #else
@ -30,7 +30,7 @@ REGISTER8(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, int8, int16, int32,
REGISTER4(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, uint8, uint16, REGISTER4(BinaryOp, GPU, "BitwiseXor", functor::bitwise_xor, uint8, uint16,
uint32, uint64); uint32, uint64);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED || #endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED // !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow } // 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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Ceil", functor::ceil, float, Eigen::half, double);
#endif #endif
#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"), Name("Conj").Device(DEVICE_GPU).TypeConstraint<Variant>("T"),
UnaryVariantOp<GPUDevice, CONJ_VARIANT_UNARY_OP>); UnaryVariantOp<GPUDevice, CONJ_VARIANT_UNARY_OP>);
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \
!defined(MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED) !defined(MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED)
REGISTER_KERNEL_BUILDER( REGISTER_KERNEL_BUILDER(
Name("Conj").Device(DEVICE_GPU).TypeConstraint<complex64>("T"), Name("Conj").Device(DEVICE_GPU).TypeConstraint<complex64>("T"),
UnaryOp<GPUDevice, functor::conj<complex64>>); 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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Cos", functor::cos, float, Eigen::half, double);
#endif #endif
#endif #endif

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

@ -28,7 +28,7 @@ REGISTER_KERNEL_BUILDER(
ApproximateEqualOp<CPUDevice, double>); ApproximateEqualOp<CPUDevice, double>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER4(BinaryOp, GPU, "Equal", functor::equal_to, float, Eigen::half, double,
uint8); uint8);
#else #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); complex128, tstring, bool);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER6(BinaryOp, GPU, "Equal", functor::equal_to, int8, int16, int64,
complex64, complex128, bool); complex64, complex128, bool);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER5(UnaryOp, GPU, "Exp", functor::exp, float, Eigen::half, double,
complex64, complex128); complex64, complex128);
#endif #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Floor", functor::floor, float, Eigen::half, double);
#endif #endif
#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 tensorflow {
namespace functor { namespace functor {
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, DEFINE_BINARY8(left_shift, int8, int16, int32, int64, uint8, uint16, uint32,
uint64); uint64);
#else #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 tensorflow {
namespace functor { namespace functor {
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, DEFINE_BINARY8(right_shift, int8, int16, int32, int64, uint8, uint16, uint32,
uint64); uint64);
#else #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); double, int32, int64, uint8, int8, int16, bfloat16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER7(BinaryOp, GPU, "Greater", functor::greater, float, Eigen::half,
double, int64, uint8, int8, int16); double, int64, uint8, int8, int16);
#else #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); Eigen::half, double, int32, int64, uint8, int8, int16, bfloat16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER7(BinaryOp, GPU, "GreaterEqual", functor::greater_equal, float,
Eigen::half, double, int64, uint8, int8, int16); Eigen::half, double, int64, uint8, int8, int16);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, float, complex64);
REGISTER_COMPLEX(GPU, double, complex128); REGISTER_COMPLEX(GPU, double, complex128);
#endif #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "IsInf", functor::isinf, float, Eigen::half, double);
#endif #endif
#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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER8(BinaryOp, GPU, "LeftShift", functor::left_shift, int8, int16, int32,
int64, uint8, uint16, uint32, uint64); int64, uint8, uint16, uint32, uint64);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER7(BinaryOp, GPU, "Less", functor::less, float, Eigen::half, double,
int64, uint8, int8, int16); int64, uint8, int8, int16);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER7(BinaryOp, GPU, "LessEqual", functor::less_equal, float, Eigen::half,
double, int64, uint8, int8, int16); double, int64, uint8, int8, int16);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Log", functor::log, float, Eigen::half, double);
#endif #endif
#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>); BinaryOp<CPUDevice, functor::logical_and>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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), REGISTER_KERNEL_BUILDER(Name("LogicalAnd").Device(DEVICE_GPU),
BinaryOp<GPUDevice, functor::logical_and>); BinaryOp<GPUDevice, functor::logical_and>);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED || #endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED // !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif #endif
} // namespace tensorflow } // 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>); UnaryOp<CPUDevice, functor::logical_not>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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), REGISTER_KERNEL_BUILDER(Name("LogicalNot").Device(DEVICE_GPU),
UnaryOp<GPUDevice, functor::logical_not>); UnaryOp<GPUDevice, functor::logical_not>);
#endif #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>); BinaryOp<CPUDevice, functor::logical_or>);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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), REGISTER_KERNEL_BUILDER(Name("LogicalOr").Device(DEVICE_GPU),
BinaryOp<GPUDevice, functor::logical_or>); BinaryOp<GPUDevice, functor::logical_or>);
#endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED || #endif // !MLIR_GENERATED_GPU_KERNELS_ENABLED ||
// !MLIR_GENERATED_UNRANKED_GPU_KERNELS_ENABLED // !MLIR_GENERATED_EXPERIMENTAL_GPU_KERNELS_ENABLED
#endif #endif
} // namespace tensorflow } // 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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER4(BinaryOp, GPU, "Mul", functor::mul, Eigen::half, float, double,
uint8); uint8);
#else #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); complex64, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER6(BinaryOp, GPU, "Mul", functor::mul, int8, uint16, int16, int64,
complex64, complex128); complex64, complex128);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Neg", functor::neg, int8, int16, int64);
#endif #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER6(UnaryOp, GPU, "Neg", functor::neg, Eigen::half, float, double,
bfloat16, complex64, complex128); bfloat16, complex64, complex128);
#else #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); uint64, qint8, qint16, quint8, quint16);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER4(BinaryOp, GPU, "NotEqual", functor::not_equal_to, float, Eigen::half,
double, uint8); double, uint8);
#else #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); complex64, complex128, tstring, bool);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER6(BinaryOp, GPU, "NotEqual", functor::not_equal_to, int8, int16, int64,
complex64, complex128, bool); complex64, complex128, bool);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, float, complex64);
REGISTER_COMPLEX(GPU, double, complex128); REGISTER_COMPLEX(GPU, double, complex128);
#endif #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER8(BinaryOp, GPU, "RightShift", functor::right_shift, int8, int16, int32,
int64, uint8, uint16, uint32, uint64); int64, uint8, uint16, uint32, uint64);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Rsqrt", functor::rsqrt, float, Eigen::half, double);
#endif #endif
#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); complex64, Eigen::half, bfloat16, complex128);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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, REGISTER6(UnaryOp, GPU, "Sign", functor::sign, float, Eigen::half, double,
int64, complex64, complex128); int64, complex64, complex128);
#else #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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Sin", functor::sin, float, Eigen::half, double);
#endif #endif
#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 GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) || \ #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); REGISTER3(UnaryOp, GPU, "Sqrt", functor::sqrt, float, Eigen::half, double);
#endif #endif
#endif #endif

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

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