Skip to content

[seq2seq] fix gpu test #350

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 8, 2019
Merged

[seq2seq] fix gpu test #350

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
bad20b8
fix gpu test
WindQAQ Jul 13, 2019
635e13b
Merge branch 'master' of https://github.com/WindQAQ/addons into seq2s…
WindQAQ Aug 2, 2019
c7513d4
build gpu kernel for beam_search_ops
WindQAQ Aug 4, 2019
f006c9b
remove assertion manager
WindQAQ Aug 4, 2019
fc3c6ff
Merge branch 'master' of https://github.com/WindQAQ/addons into seq2s…
WindQAQ Aug 4, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions tensorflow_addons/custom_ops/seq2seq/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,23 +3,45 @@ 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")

cc_binary(
name = "_beam_search_ops.so",
srcs = [
"cc/kernels/beam_search_ops.cc",
"cc/kernels/beam_search_ops.h",
# "cc/kernels/beam_search_ops_gpu.cu.cc",
"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 = [
"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_libs",
"@local_config_cuda//cuda:cuda_headers",
]),
alwayslink = 1,
)
2 changes: 1 addition & 1 deletion tensorflow_addons/custom_ops/seq2seq/cc/kernels/beam_search_ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,4 +37,4 @@ struct GatherTree {
} // namespace functor
} // namespace tensorflow

#endif // TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_
#endif // TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_
2 changes: 1 addition & 1 deletion tensorflow_addons/custom_ops/seq2seq/cc/kernels/beam_search_ops_gpu.cu.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ limitations under the License.
#define EIGEN_USE_GPU

#include "tensorflow_addons/custom_ops/seq2seq/cc/kernels/beam_search_ops.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"

namespace tensorflow {
namespace functor {
Expand Down
3 changes: 0 additions & 3 deletions tensorflow_addons/seq2seq/beam_search_ops_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,6 @@ def testBadParentValuesOnCPU(self):
self.evaluate(beams)

def testBadParentValuesOnGPU(self):
# TODO: Fix #348 issue
self.skipTest('Wait #348 to be fixed')

# Only want to run this test on CUDA devices, as gather_tree is not
# registered for SYCL devices.
if not tf.test.is_gpu_available(cuda_only=True):
Expand Down