Internal Change.

Change: 154605009

1 files changed, 51 insertions, 32 deletions

diff --git a/tensorflow/contrib/nccl/BUILD b/tensorflow/contrib/nccl/BUILD
index c7f32baa2d..b6314079e1 100644
--- a/tensorflow/contrib/nccl/BUILD
+++ b/tensorflow/contrib/nccl/BUILD
@@ -1,10 +1,8 @@
 # Description:
 #   Wrap NVIDIA (https://github.com/NVIDIA/nccl) NCCL with tensorflow ops.
 #   APIs are meant to change over time.
-package(
-    default_visibility = ["//visibility:private"],
-    features = ["-parse_headers"],
-)
+
+package(default_visibility = ["//tensorflow:__subpackages__"])
 
 licenses(["notice"])  # Apache 2.0
 
@@ -19,6 +17,8 @@ load(
 )
 load("//tensorflow:tensorflow.bzl", "cuda_py_test")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
+load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
 
 tf_custom_op_library(
     name = "python/ops/_nccl_ops.so",
@@ -34,6 +34,47 @@ tf_custom_op_library(
     ],
 )
 
+tf_cuda_cc_test(
+    name = "nccl_manager_test",
+    size = "medium",
+    srcs = if_cuda(
+        [
+            "kernels/nccl_manager.cc",
+            "kernels/nccl_manager.h",
+            "kernels/nccl_manager_test.cc",
+        ],
+        [],
+    ),
+    tags = ["manual"],  # Disabled until errors finding nvmlShutdown are found.
+    deps = if_cuda(
+        [
+            "@nccl_archive//:nccl",
+            "//tensorflow/core",
+            "//tensorflow/core:cuda",
+        ],
+        [],
+    ) + [
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+tf_kernel_library(
+    name = "nccl_kernels",
+    srcs = [
+        "kernels/nccl_manager.cc",
+        "kernels/nccl_manager.h",
+        "kernels/nccl_ops.cc",
+    ],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:gpu_headers_lib",
+        "@nccl_archive//:nccl",
+    ],
+    alwayslink = 1,
+)
+
 tf_gen_op_libs(
     op_lib_names = ["nccl_ops"],
     deps = [
@@ -46,15 +87,19 @@ tf_gen_op_wrapper_py(
     deps = [":nccl_ops_op_lib"],
 )
 
-py_library(
+tf_custom_op_py_library(
     name = "nccl_py",
     srcs = [
         "__init__.py",
         "python/ops/nccl_ops.py",
     ],
-    data = [
+    dso = [
         ":python/ops/_nccl_ops.so",
     ],
+    kernels = [
+        ":nccl_kernels",
+        ":nccl_ops_op_lib",
+    ],
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
     deps = [
@@ -82,32 +127,6 @@ cuda_py_test(
     ],
 )
 
-tf_cuda_cc_test(
-    name = "nccl_manager_test",
-    size = "medium",
-    srcs = if_cuda(
-        [
-            "kernels/nccl_manager.cc",
-            "kernels/nccl_manager.h",
-            "kernels/nccl_manager_test.cc",
-        ],
-        [],
-    ),
-    tags = ["manual"],  # Disabled until errors finding nvmlShutdown are found.
-    deps = if_cuda(
-        [
-            "@nccl_archive//:nccl",
-            "//tensorflow/core",
-            "//tensorflow/core:cuda",
-        ],
-        [],
-    ) + [
-        "//tensorflow/core:test",
-        "//tensorflow/core:test_main",
-        "//tensorflow/core:testlib",
-    ],
-)
-
 filegroup(
     name = "all_files",
     srcs = glob(