Enforce the rule that custom op libraries don't depend on core:framew…

…ork or core:lib Change: 122300990
tensorflow · May 13, 2016 · febd091 · febd091
1 parent 0ffcb6c
commit febd091
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 4 deletions.
diff --git a/tensorflow/g3doc/how_tos/adding_an_op/index.md b/tensorflow/g3doc/how_tos/adding_an_op/index.md
@@ -172,6 +172,12 @@ Run the following command to build `zero_out.so`.
 $ bazel build -c opt //tensorflow/core/user_ops:zero_out.so
 ```
 
+> Note:
+Although you can create a shared library (a `.so` file) with the standard
+`cc_library` rule, we strongly recommend that you use the `tf_custom_op_library`
+macro. It adds some required dependencies, and performs checks to ensure that
+the shared library is compatible with TensorFlow's plugin loading mechanism.
+
 ## Using the Op in Python
 
 TensorFlow Python API provides the

diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
@@ -520,26 +520,61 @@ def tf_custom_op_library_additional_deps():
       "//tensorflow/core:framework_headers_lib",
   ]
 
+# Process the "deps" attribute of the rule.
+# Raise an error if deps has one of the "disallowed_deps".
+# Note : The aspect function gets called on all targets of the dependency
+# graph reachable via "deps" edges.
+def _collect_deps_aspect_impl(target, ctx):
+  disallowed_deps = ["tensorflow/core:framework", "tensorflow/core:lib"]
+  if hasattr(ctx.rule.attr, "deps"):
+    for dep in ctx.rule.attr.deps:
+      label_name = dep.label.package + ":" + dep.label.name
+      for disallowed_dep in disallowed_deps:
+        if label_name.endswith(disallowed_dep):
+          fail("tf_custom_op_library cannot depend on " + disallowed_dep)
+  return struct()
+
+collect_deps_aspect = aspect(
+    implementation=_collect_deps_aspect_impl,
+    attr_aspects=["deps"])
+
+# A dummy rule just so we can have an aspect function process the transitive
+# dependencies of a cc_{library, binary}
+def _collect_deps_impl(ctx):
+  return struct()
+
+collect_deps = rule(
+    _collect_deps_impl,
+    attrs = {
+        "deps": attr.label_list(
+            aspects=[collect_deps_aspect],
+        )},
+)
+
 # Helper to build a dynamic library (.so) from the sources containing
 # implementations of custom ops and kernels.
 def tf_custom_op_library(name, srcs=[], gpu_srcs=[], deps=[]):
   cuda_deps = [
       "//tensorflow/core:stream_executor_headers_lib",
       "//third_party/gpus/cuda:cudart_static",
   ]
-  deps = deps + tf_custom_op_library_additional_deps()
+  deps = deps + tf_custom_op_library_additional_deps() + if_cuda(cuda_deps)
   if gpu_srcs:
     basename = name.split(".")[0]
     native.cc_library(
         name = basename + "_gpu",
         srcs = gpu_srcs,
         copts = _cuda_copts(),
-        deps = deps + if_cuda(cuda_deps))
-    cuda_deps.extend([":" + basename + "_gpu"])
+        deps = deps)
+    deps.extend([":" + basename + "_gpu"])
+
+  collect_deps(name=name+"_check_deps",
+               deps=deps)
 
   native.cc_binary(name=name,
                    srcs=srcs,
-                   deps=deps + if_cuda(cuda_deps),
+                   deps=deps,
+                   data=[name + "_check_deps"],
                    linkshared=1,
                    linkopts = select({
                        "//conditions:default": [