diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-04-05 03:09:27 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-05 03:11:33 -0700 |
| commit | 1fda7645d132b71b9084b01945795e97e582adcd (patch) | |
| tree | 7ae6dfa150c4b58dc339db0ae1f0b177c482433e /tensorflow/contrib/nccl | |
| parent | dfa9921e6343727b05f42f8d4a918b19528ff994 (diff) | |
Add support for NCCL2. The configure script asks for what version of NCCL to use. The default is still NCCL 1 from GitHub. If the user chooses NCCL 2, it asks for the install directory.
The nccl_configure.bzl generates two different BUILD files based on the chose NCCL version. For NCCL 1, it aliases to the existing 'nccl_archive' http_repo on GitHub. For NCCL 2, it creates a target containing the NCCL 2 library and headers from the chosen install directory.
PiperOrigin-RevId: 191718007
Diffstat (limited to 'tensorflow/contrib/nccl')
| -rw-r--r-- | tensorflow/contrib/nccl/BUILD | 6 | ||||
| -rw-r--r-- | tensorflow/contrib/nccl/kernels/nccl_manager.h | 2 | ||||
| -rw-r--r-- | tensorflow/contrib/nccl/kernels/nccl_ops.cc | 2 |
3 files changed, 5 insertions, 5 deletions
diff --git a/tensorflow/contrib/nccl/BUILD b/tensorflow/contrib/nccl/BUILD index 6cbfd03881..334e70318d 100644 --- a/tensorflow/contrib/nccl/BUILD +++ b/tensorflow/contrib/nccl/BUILD @@ -31,7 +31,7 @@ tf_custom_op_library( "kernels/nccl_ops.cc", ], deps = if_cuda([ - "@nccl_archive//:nccl", + "@local_config_nccl//:nccl", "//tensorflow/core:gpu_headers_lib", ]), ) @@ -61,7 +61,7 @@ tf_cuda_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", - "@nccl_archive//:nccl", + "@local_config_nccl//:nccl", ], ) @@ -80,7 +80,7 @@ tf_kernel_library( "//tensorflow/core:lib", "//tensorflow/core:proto_text", "//tensorflow/core:stream_executor", - "@nccl_archive//:nccl", + "@local_config_nccl//:nccl", ], alwayslink = 1, ) diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager.h b/tensorflow/contrib/nccl/kernels/nccl_manager.h index bb219e0edc..6ff8cea84e 100644 --- a/tensorflow/contrib/nccl/kernels/nccl_manager.h +++ b/tensorflow/contrib/nccl/kernels/nccl_manager.h @@ -20,7 +20,7 @@ limitations under the License. #include <unordered_map> #include <vector> -#include "src/nccl.h" +#include "third_party/nccl/nccl.h" #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/platform/mutex.h" diff --git a/tensorflow/contrib/nccl/kernels/nccl_ops.cc b/tensorflow/contrib/nccl/kernels/nccl_ops.cc index 266d4f6f0d..c2b76caef3 100644 --- a/tensorflow/contrib/nccl/kernels/nccl_ops.cc +++ b/tensorflow/contrib/nccl/kernels/nccl_ops.cc @@ -17,7 +17,7 @@ limitations under the License. #include <vector> -#include "src/nccl.h" +#include "third_party/nccl/nccl.h" #include "tensorflow/contrib/nccl/kernels/nccl_manager.h" #include "tensorflow/core/framework/op_kernel.h" |