diff --git a/tensorflow_addons/custom_ops/activations/BUILD b/tensorflow_addons/custom_ops/activations/BUILD index 56b0f00d3b..f0b7624b4a 100644 --- a/tensorflow_addons/custom_ops/activations/BUILD +++ b/tensorflow_addons/custom_ops/activations/BUILD @@ -2,120 +2,9 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//visibility:public"]) -load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda") +load("//tensorflow_addons:tensorflow_addons.bzl", "custom_op_library") -cc_library( - name = "gelu_op_gpu", - srcs = [ - "cc/kernels/gelu_op.h", - "cc/kernels/gelu_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, -) - -cc_library( - name = "hardshrink_op_gpu", - srcs = [ - "cc/kernels/hardshrink_op.h", - "cc/kernels/hardshrink_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, -) - -cc_library( - name = "lisht_op_gpu", - srcs = [ - "cc/kernels/lisht_op.h", - "cc/kernels/lisht_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_libs", - "@local_config_cuda//cuda:cuda_headers", - ]), - alwayslink = 1, -) - -cc_library( - name = "softshrink_op_gpu", - srcs = [ - "cc/kernels/softshrink_op.h", - "cc/kernels/softshrink_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, -) - -cc_library( - name = "tanhshrink_op_gpu", - srcs = [ - "cc/kernels/tanhshrink_op.h", - "cc/kernels/tanhshrink_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, -) - -cc_binary( +custom_op_library( name = "_activation_ops.so", srcs = [ "cc/kernels/gelu_op.cc", @@ -134,20 +23,16 @@ cc_binary( "cc/ops/softshrink_op.cc", "cc/ops/tanhshrink_op.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ] + if_cuda(["-DGOOGLE_CUDA=1"]), - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - ":gelu_op_gpu", - ":hardshrink_op_gpu", - ":lisht_op_gpu", - ":softshrink_op_gpu", - ":tanhshrink_op_gpu", - ]), + cuda_srcs = [ + "cc/kernels/gelu_op.h", + "cc/kernels/gelu_op_gpu.cu.cc", + "cc/kernels/hardshrink_op.h", + "cc/kernels/hardshrink_op_gpu.cu.cc", + "cc/kernels/lisht_op.h", + "cc/kernels/lisht_op_gpu.cu.cc", + "cc/kernels/softshrink_op.h", + "cc/kernels/softshrink_op_gpu.cu.cc", + "cc/kernels/tanhshrink_op.h", + "cc/kernels/tanhshrink_op_gpu.cu.cc", + ], ) diff --git a/tensorflow_addons/custom_ops/image/BUILD b/tensorflow_addons/custom_ops/image/BUILD index bfc0054835..7002d86e4a 100644 --- a/tensorflow_addons/custom_ops/image/BUILD +++ b/tensorflow_addons/custom_ops/image/BUILD @@ -2,95 +2,22 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//visibility:public"]) -load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda") +load("//tensorflow_addons:tensorflow_addons.bzl", "custom_op_library") -cc_library( - name = "distort_image_ops_gpu", - srcs = [ - "cc/kernels/adjust_hsv_in_yiq_op.h", - "cc/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, -) - -cc_binary( +custom_op_library( name = "_distort_image_ops.so", srcs = [ "cc/kernels/adjust_hsv_in_yiq_op.cc", "cc/kernels/adjust_hsv_in_yiq_op.h", "cc/ops/distort_image_ops.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ] + if_cuda(["-DGOOGLE_CUDA=1"]), - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([":distort_image_ops_gpu"]), -) - -cc_library( - name = "image_projective_transform_op_gpu", - srcs = [ - "cc/kernels/image_projective_transform_op.h", - "cc/kernels/image_projective_transform_op_gpu.cu.cc", - ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, -) - -cc_library( - name = "euclidean_distance_transform_op_gpu", - srcs = [ - "cc/kernels/euclidean_distance_transform_op.h", - "cc/kernels/euclidean_distance_transform_op_gpu.cu.cc", + cuda_srcs = [ + "cc/kernels/adjust_hsv_in_yiq_op.h", + "cc/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc", ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, ) -cc_binary( +custom_op_library( name = "_image_ops.so", srcs = [ "cc/kernels/connected_components.cc", @@ -101,17 +28,10 @@ cc_binary( "cc/kernels/image_projective_transform_op.h", "cc/ops/image_ops.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ] + if_cuda(["-DGOOGLE_CUDA=1"]), - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - ":image_projective_transform_op_gpu", - ":euclidean_distance_transform_op_gpu", - ]), + cuda_srcs = [ + "cc/kernels/euclidean_distance_transform_op.h", + "cc/kernels/euclidean_distance_transform_op_gpu.cu.cc", + "cc/kernels/image_projective_transform_op.h", + "cc/kernels/image_projective_transform_op_gpu.cu.cc", + ], ) diff --git a/tensorflow_addons/custom_ops/layers/BUILD b/tensorflow_addons/custom_ops/layers/BUILD index 9ad39a5f4e..5ff49efe99 100644 --- a/tensorflow_addons/custom_ops/layers/BUILD +++ b/tensorflow_addons/custom_ops/layers/BUILD @@ -2,47 +2,20 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//visibility:public"]) -load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda") +load("//tensorflow_addons:tensorflow_addons.bzl", "custom_op_library") -cc_binary( +custom_op_library( name = "_correlation_cost_ops.so", srcs = [ "cc/kernels/correlation_cost_op.cc", "cc/kernels/correlation_cost_op.h", "cc/ops/correlation_cost_op.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ] + if_cuda(["-DGOOGLE_CUDA=1"]), - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([":correlation_cost_ops_gpu"]), -) - -cc_library( - name = "correlation_cost_ops_gpu", - srcs = [ + cuda_deps = [ + "@cub_archive//:cub", + ], + cuda_srcs = [ "cc/kernels/correlation_cost_op.h", "cc/kernels/correlation_cost_op_gpu.cu.cc", ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - "@cub_archive//:cub", - ]), - alwayslink = 1, ) diff --git a/tensorflow_addons/custom_ops/seq2seq/BUILD b/tensorflow_addons/custom_ops/seq2seq/BUILD index 712c78e1c4..9abe889530 100644 --- a/tensorflow_addons/custom_ops/seq2seq/BUILD +++ b/tensorflow_addons/custom_ops/seq2seq/BUILD @@ -2,46 +2,17 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//visibility:public"]) -load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda") +load("//tensorflow_addons:tensorflow_addons.bzl", "custom_op_library") -cc_binary( +custom_op_library( name = "_beam_search_ops.so", srcs = [ "cc/kernels/beam_search_ops.cc", "cc/kernels/beam_search_ops.h", "cc/ops/beam_search_ops.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ] + if_cuda(["-DGOOGLE_CUDA=1"]), - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([":beam_search_ops_gpu"]), -) - -cc_library( - name = "beam_search_ops_gpu", - srcs = [ + cuda_srcs = [ "cc/kernels/beam_search_ops.h", "cc/kernels/beam_search_ops_gpu.cu.cc", ], - copts = if_cuda_is_configured([ - "-DGOOGLE_CUDA=1", - "-x cuda", - "-nvcc_options=relaxed-constexpr", - "-nvcc_options=ftz=true", - ]), - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ] + if_cuda_is_configured([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:cudart_static", - ]), - alwayslink = 1, ) diff --git a/tensorflow_addons/custom_ops/text/BUILD b/tensorflow_addons/custom_ops/text/BUILD index 3f2a03f3db..a52e139228 100644 --- a/tensorflow_addons/custom_ops/text/BUILD +++ b/tensorflow_addons/custom_ops/text/BUILD @@ -2,40 +2,20 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//visibility:public"]) -load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI") +load("//tensorflow_addons:tensorflow_addons.bzl", "custom_op_library") -cc_binary( +custom_op_library( name = "_skip_gram_ops.so", srcs = [ "cc/kernels/skip_gram_kernels.cc", "cc/ops/skip_gram_ops.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ], - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ], ) -cc_binary( +custom_op_library( name = "_parse_time_op.so", srcs = [ "cc/kernels/parse_time_kernel.cc", "cc/ops/parse_time_op.cc", ], - copts = [ - "-pthread", - "-std=c++11", - D_GLIBCXX_USE_CXX11_ABI, - ], - linkshared = 1, - deps = [ - "@local_config_tf//:libtensorflow_framework", - "@local_config_tf//:tf_header_lib", - ], ) diff --git a/tensorflow_addons/tensorflow_addons.bzl b/tensorflow_addons/tensorflow_addons.bzl new file mode 100644 index 0000000000..2bdee54e52 --- /dev/null +++ b/tensorflow_addons/tensorflow_addons.bzl @@ -0,0 +1,52 @@ +load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI") +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda") + +def custom_op_library( + name, + srcs = [], + cuda_srcs = [], + deps = [], + cuda_deps = [], + copts = [], + **kwargs): + deps = deps + [ + "@local_config_tf//:libtensorflow_framework", + "@local_config_tf//:tf_header_lib", + ] + + if cuda_srcs: + copts = copts + if_cuda(["-DGOOGLE_CUDA=1"]) + cuda_copts = copts + if_cuda_is_configured([ + "-x cuda", + "-nvcc_options=relaxed-constexpr", + "-nvcc_options=ftz=true", + ]) + cuda_deps = deps + if_cuda_is_configured(cuda_deps) + if_cuda_is_configured([ + "@local_config_cuda//cuda:cuda_headers", + "@local_config_cuda//cuda:cudart_static", + ]) + basename = name.split(".")[0] + native.cc_library( + name = basename + "_gpu", + srcs = cuda_srcs, + deps = cuda_deps, + copts = cuda_copts, + alwayslink = 1, + **kwargs + ) + deps = deps + if_cuda_is_configured([":" + basename + "_gpu"]) + + copts = copts + [ + "-pthread", + "-std=c++11", + D_GLIBCXX_USE_CXX11_ABI, + ] + + native.cc_binary( + name = name, + srcs = srcs, + copts = copts, + linkshared = 1, + deps = deps, + **kwargs + )