Untangle internal build python and cpp dependencies

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/18326

Reviewed By: ezyang

Differential Revision: D14576294

fbshipit-source-id: 186ce1e3d026d962b7386f861eddf093f583a878

diff --git a/tools/build_variables.py b/tools/build_variables.py
index 686287b..675ae7d 100644 (file)
--- a/tools/build_variables.py
+++ b/tools/build_variables.py
@@ -117,6 +117,7 @@ libtorch_sources = [
     "torch/csrc/jit/fuser/cpu/fused_kernel.cpp",
     "torch/csrc/jit/fuser/cpu/dynamic_library_unix.cpp",
     "torch/csrc/jit/fuser/interface.cpp",
+    "test/cpp/jit/test.cpp",
 ]
 
 libtorch_cuda_sources = [
@@ -141,6 +142,8 @@ def add_torch_libs():
             # distributed only uses Module.cpp
             # so remove all other files and just include that
             "torch/csrc/distributed/**/*.cpp",
+            # top-level hook of extension registration lives in a separate file
+            "torch/csrc/stub.cpp",
         ],
     ) + [
         "torch/csrc/distributed/Module.cpp",
@@ -179,7 +182,7 @@ def add_torch_libs():
                 "-Wno-unknown-pragmas",
             ]
         },
-        "headers": native.glob(["torch/csrc/**/*.h", "torch/csrc/generic/*.cpp"]),
+        "headers": native.glob(["torch/csrc/**/*.h", "torch/csrc/generic/*.cpp", "test/cpp/jit/*.h"]),
         "preprocessor_flags": [
             "-Icaffe2",
             "-Icaffe2/torch/csrc/api/include",
@@ -238,14 +241,13 @@ def add_torch_libs():
         **common_flags
     )
 
-    cpp_python_extension(
-        name="_C",
+    # TODO: split it into cpp and cuda parts similarly to libtorch
+    cpp_library(
+        name="_C_impl",
         srcs=libtorch_python_sources,
-        base_module="torch",
         deps=[
             ":libtorch_cuda",
             ":thnn",
-            ":torch-lib-headers",
             "//caffe2/torch/lib/THD:THD",
             "//caffe2/torch/lib/c10d:c10d",
             "//caffe2/torch/lib/libshm:libshm",
@@ -253,8 +255,18 @@ def add_torch_libs():
         external_deps=[
             ("numpy", None, "cpp"),
             ("pybind11", None),
+            ("python", None),
         ],
         **common_flags
     )
 
+    cpp_python_extension(
+        name="_C",
+        srcs=[
+            "torch/csrc/stub.cpp",
+        ],
+        base_module="torch",
+        deps=[":_C_impl"],
+    )
+
     return r