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

Split the unary unranked kernels into enabled by default and experimental.

Browse Source 
For now, this change will be a no-op, but in a followup change we will
redefine the meaning of --define=enable_unranked_kernels=1 to enable the
experimental unary kernels, while the others (currently abs and tanh) will be
enabled by default and can be disabled with
--define=tensorflow_enable_mlir_generated_gpu_kernels=0
The ranked unary kernels will be removed.
PiperOrigin-RevId: 345419524
Change-Id: Ice00e32f695fd44d1b70b8cf8c377f6a61de3b19
This commit is contained in:
Adrian Kuegel 2020-12-03 03:46:33 -08:00 committed by TensorFlower Gardener
parent 7fa02aac5f
commit 6c0245b27c
1 changed files with 46 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -38,16 +38,17 @@ config_setting(
)
filegroup(
name = "unary_kernel_srcs",
srcs = if_mlir_unranked_kernels_enabled(
if_false = [
"cwise_op_gpu_abs.cc",
"cwise_op_gpu_base.cc",
"cwise_op_gpu_base.h",
"cwise_op_gpu_tanh.cc",
],
if_true = [
name = "enabled_unary_unranked_kernel_srcs",
srcs = [
"unranked_op_gpu_abs.cc",
"unranked_op_gpu_tanh.cc",
],
compatible_with = get_compatible_with_cloud(),
)
filegroup(
name = "experimental_unary_unranked_kernel_srcs",
srcs = [
"unranked_op_gpu_ceil.cc",
"unranked_op_gpu_conj.cc",
"unranked_op_gpu_cos.cc",
@ -62,8 +63,30 @@ filegroup(
"unranked_op_gpu_sign.cc",
"unranked_op_gpu_sin.cc",
"unranked_op_gpu_sqrt.cc",
"unranked_op_gpu_tanh.cc",
],
compatible_with = get_compatible_with_cloud(),
)
filegroup(
name = "unary_unranked_kernel_srcs",
srcs = [
":enabled_unary_unranked_kernel_srcs",
] + if_mlir_unranked_kernels_enabled(
if_true = [":experimental_unary_unranked_kernel_srcs"],
),
compatible_with = get_compatible_with_cloud(),
)
filegroup(
name = "unary_kernel_srcs",
srcs = if_mlir_unranked_kernels_enabled(
if_false = [
"cwise_op_gpu_abs.cc",
"cwise_op_gpu_base.cc",
"cwise_op_gpu_base.h",
"cwise_op_gpu_tanh.cc",
],
if_true = [":unary_unranked_kernel_srcs"],
),
compatible_with = get_compatible_with_cloud(),
)
@ -104,6 +127,11 @@ tf_kernel_library(
"//tensorflow/core/platform:stream_executor",
],
if_true = [
# Technically we only need to depend on the kernel libraries for the
# unranked kernels which are enabled by default. But this would
# make our BUILD target structure uglier. We already need to make
# sure that those targets can be built, so it should not hurt to
# link them in even if they are currently not needed yet.
":abs_unranked_kernels",
":ceil_unranked_kernels",
":conj_unranked_kernels",