Caffe2 - Python API
A deep learning, cross platform ML framework
build_variables.py
1 # In the open-source build, these are generated into
2 # torch/csrc/{autgrad,jit}/generated. In fbcode, this distinction is
3 # not currently relevant so they are combined into one list.
4 from __future__ import absolute_import, division, print_function, unicode_literals
5 load("@bazel_skylib//lib:new_sets.bzl", "sets")
6 
7 
8 GENERATED_CPP = [
9  "Functions.cpp",
10  "THCUNN.cpp",
11  "THNN.cpp",
12  "VariableType_0.cpp",
13  "VariableType_1.cpp",
14  "VariableType_2.cpp",
15  "VariableType_3.cpp",
16  "VariableType_4.cpp",
17  "register_aten_ops_0.cpp",
18  "register_aten_ops_1.cpp",
19  "register_aten_ops_2.cpp",
20  "python_functions.cpp",
21  "python_nn_functions.cpp",
22  "python_torch_functions.cpp",
23  "python_variable_methods.cpp",
24 ]
25 
26 # copied from https://github.com/pytorch/pytorch/blob/master/tools/cpp_build/torch/CMakeLists.txt
27 libtorch_sources = [
28  ":generate-code=Functions.cpp",
29  ":generate-code=register_aten_ops_0.cpp",
30  ":generate-code=register_aten_ops_1.cpp",
31  ":generate-code=register_aten_ops_2.cpp",
32  ":generate-code=VariableType_0.cpp",
33  ":generate-code=VariableType_1.cpp",
34  ":generate-code=VariableType_2.cpp",
35  ":generate-code=VariableType_3.cpp",
36  ":generate-code=VariableType_4.cpp",
37  "torch/csrc/autograd/VariableTypeManual.cpp",
38  "torch/csrc/autograd/anomaly_mode.cpp",
39  "torch/csrc/autograd/engine.cpp",
40  "torch/csrc/autograd/function.cpp",
41  "torch/csrc/autograd/functions/accumulate_grad.cpp",
42  "torch/csrc/autograd/functions/basic_ops.cpp",
43  "torch/csrc/autograd/functions/tensor.cpp",
44  "torch/csrc/autograd/functions/utils.cpp",
45  "torch/csrc/autograd/grad_mode.cpp",
46  "torch/csrc/autograd/input_buffer.cpp",
47  "torch/csrc/autograd/profiler.cpp",
48  "torch/csrc/autograd/saved_variable.cpp",
49  "torch/csrc/autograd/variable.cpp",
50  "torch/csrc/Exceptions.cpp",
51  "torch/csrc/jit/autodiff.cpp",
52  "torch/csrc/jit/attributes.cpp",
53  "torch/csrc/jit/constants.cpp",
54  "torch/csrc/jit/node_hashing.cpp",
55  "torch/csrc/jit/export.cpp",
56  "torch/csrc/jit/pickler.cpp",
57  "torch/csrc/jit/graph_executor.cpp",
58  "torch/csrc/jit/import.cpp",
59  "torch/csrc/jit/interpreter.cpp",
60  "torch/csrc/jit/ir.cpp",
61  "torch/csrc/jit/irparser.cpp",
62  "torch/csrc/jit/netdef_converter.cpp",
63  "torch/csrc/jit/register_c10_ops.cpp",
64  "torch/csrc/jit/symbolic_script.cpp",
65  "torch/csrc/jit/operator.cpp",
66  "torch/csrc/jit/passes/alias_analysis.cpp",
67  "torch/csrc/jit/passes/batch_mm.cpp",
68  "torch/csrc/jit/passes/canonicalize_ops.cpp",
69  "torch/csrc/jit/passes/canonicalize.cpp",
70  "torch/csrc/jit/passes/common_subexpression_elimination.cpp",
71  "torch/csrc/jit/passes/constant_propagation.cpp",
72  "torch/csrc/jit/passes/constant_pooling.cpp",
73  "torch/csrc/jit/passes/create_autodiff_subgraphs.cpp",
74  "torch/csrc/jit/passes/dead_code_elimination.cpp",
75  "torch/csrc/jit/passes/erase_number_types.cpp",
76  "torch/csrc/jit/passes/graph_fuser.cpp",
77  "torch/csrc/jit/passes/inline_autodiff_subgraphs.cpp",
78  "torch/csrc/jit/passes/inplace_check.cpp",
79  "torch/csrc/jit/passes/loop_unrolling.cpp",
80  "torch/csrc/jit/passes/lower_grad_of.cpp",
81  "torch/csrc/jit/passes/lower_tuples.cpp",
82  "torch/csrc/jit/passes/peephole.cpp",
83  "torch/csrc/jit/passes/python_print.cpp",
84  "torch/csrc/jit/passes/remove_expands.cpp",
85  "torch/csrc/jit/passes/requires_grad_analysis.cpp",
86  "torch/csrc/jit/passes/shape_analysis.cpp",
87  "torch/csrc/jit/passes/specialize_autogradzero.cpp",
88  "torch/csrc/jit/passes/utils/subgraph_utils.cpp",
89  "torch/csrc/jit/passes/utils/memory_dag.cpp",
90  "torch/csrc/jit/register_prim_ops.cpp",
91  "torch/csrc/jit/register_special_ops.cpp",
92  "torch/csrc/jit/scope.cpp",
93  "torch/csrc/jit/script/compiler.cpp",
94  "torch/csrc/jit/script/edit_distance.cpp",
95  "torch/csrc/jit/script/final_returns.cpp",
96  "torch/csrc/jit/script/schema_type_parser.cpp",
97  "torch/csrc/jit/script/script_type_parser.cpp",
98  "torch/csrc/jit/script/sugared_value.cpp",
99  "torch/csrc/jit/script/schema_matching.cpp",
100  "torch/csrc/jit/script/class_type.cpp",
101  "torch/csrc/jit/script/parser.cpp",
102  "torch/csrc/jit/testing/file_check.cpp",
103  "torch/csrc/jit/import_source.cpp",
104  "torch/csrc/jit/hooks_for_testing.cpp",
105  "torch/csrc/jit/script/builtin_functions.cpp",
106  "torch/csrc/jit/script/lexer.cpp",
107  "torch/csrc/jit/script/module.cpp",
108  "torch/csrc/jit/tracer.cpp",
109  "torch/csrc/utils/tensor_flatten.cpp",
110  "torch/csrc/utils/variadic.cpp",
111  "torch/csrc/jit/fuser/kernel_cache.cpp",
112  "torch/csrc/jit/fuser/compiler.cpp",
113  "torch/csrc/jit/fuser/executor.cpp",
114  "torch/csrc/jit/fuser/codegen.cpp",
115  "torch/csrc/jit/fuser/fallback.cpp",
116  "torch/csrc/jit/fuser/cpu/fused_kernel.cpp",
117  "torch/csrc/jit/fuser/cpu/dynamic_library_unix.cpp",
118  "torch/csrc/jit/fuser/interface.cpp",
119 ]
120 
121 libtorch_cuda_sources = [
122  "torch/csrc/cuda/comm.cpp",
123  "torch/csrc/cuda/nccl.cpp",
124  "torch/csrc/jit/fuser/cuda/fused_kernel.cpp",
125  "torch/csrc/jit/fuser/cuda/thnvrtc.cpp",
126  "torch/csrc/autograd/profiler_cuda.cpp",
127  "torch/csrc/autograd/functions/comm.cpp"
128 ]
129 
130 
131 def add_torch_libs():
132  r = {}
133  # We start torch_python_sources with all cpp files, and exclude some
134  # including the files already contained in the torch and cuda bindings
135  globbed_sources = (native.glob(
136  ["torch/csrc/**/*.cpp"],
137  exclude=[
138  # remove anything that has "generic" in it"s path
139  "torch/csrc/**/generic/**/*.cpp",
140  # distributed only uses Module.cpp
141  # so remove all other files and just include that
142  "torch/csrc/distributed/**/*.cpp",
143  ],
144  ) + [
145  "torch/csrc/distributed/Module.cpp",
146  "torch/csrc/distributed/c10d/init.cpp",
147  "torch/csrc/distributed/c10d/ddp.cpp",
148  ] + [":generate-code=" + x for x in GENERATED_CPP])
149  libtorch_python_sources = sets.to_list(sets.difference(
150  sets.make(globbed_sources),
151  sets.make(libtorch_sources + libtorch_cuda_sources),
152  ))
153 
154  common_flags = {
155  "compiler_flags": [
156  "-D_THP_CORE",
157  "-DUSE_C10D",
158  "-DUSE_CUDNN",
159  "-DUSE_DISTRIBUTED",
160  "-DUSE_NCCL",
161  "-DUSE_NUMPY",
162  "-DUSE_SCALARS",
163  "-DNO_CUDNN_DESTROY_HANDLE",
164  "-DPYTORCH_ONNX_CAFFE2_BUNDLE",
165  "-Wno-write-strings",
166  "-Wno-format",
167  "-Wno-strict-aliasing",
168  "-Wno-non-virtual-dtor",
169  "-Wno-shadow-compatible-local",
170  "-Wno-empty-body",
171  ],
172  "compiler_specific_flags": {
173  "clang": [
174  "-Wno-absolute-value",
175  "-Wno-expansion-to-defined",
176  "-Wno-pessimizing-move",
177  "-Wno-return-type-c-linkage",
178  "-Wno-unknown-pragmas",
179  ]
180  },
181  "headers": native.glob(["torch/csrc/**/*.h", "torch/csrc/generic/*.cpp"]),
182  "preprocessor_flags": [
183  "-Icaffe2",
184  "-Icaffe2/torch/csrc/api/include",
185  "-Icaffe2/torch/csrc",
186  "-Icaffe2/torch/csrc/nn",
187  "-Icaffe2/torch/lib",
188  ],
189  }
190 
191  cpp_library(
192  name="libtorch",
193  srcs=libtorch_sources,
194  link_whole=True,
195  deps=[
196  ":generated-autograd-headers",
197  ":generated-autograd-headers-bare",
198  ":generated-jit-headers",
199  "//caffe2/aten:ATen-cpu",
200  "//caffe2/caffe2:caffe2_cpu",
201  "//caffe2/torch/lib/libshm:libshm",
202  ],
203  external_deps=[
204  ("nanopb", None, "protobuf-nanopb"),
205  ("protobuf", None),
206  ],
207  **common_flags
208  )
209 
210  cpp_library(
211  name="libtorch_cuda",
212  srcs=libtorch_cuda_sources,
213  link_whole=True,
214  propagated_pp_flags=[
215  "-DUSE_CUDA",
216  "-DUSE_DIRECT_NVRTC",
217  ],
218  deps=[
219  ":generated-autograd-headers",
220  ":generated-autograd-headers-bare",
221  ":generated-jit-headers",
222  ":libtorch",
223  "//caffe2/aten:ATen",
224  "//caffe2/aten:generated-aten-headers-cuda",
225  "//caffe2/caffe2:caffe2_cpu",
226  "//caffe2/caffe2:caffe2_gpu",
227  "//caffe2/torch/lib/libshm:libshm",
228  ],
229  external_deps=[
230  ("cudnn", "7.1.2", "cudnn-lazy"),
231  ("nccl", "2.1.15", "nccl-lazy"),
232  ("cuda", None, "nvToolsExt-lazy"),
233  ("cuda", None, "nvrtc-lazy"),
234  ("cuda", None, "nvrtc-builtins-lazy"),
235  ],
236  **common_flags
237  )
238 
239  cpp_python_extension(
240  name="_C",
241  srcs=libtorch_python_sources,
242  base_module="torch",
243  deps=[
244  ":libtorch_cuda",
245  ":thnn",
246  ":torch-lib-headers",
247  "//caffe2/torch/lib/THD:THD",
248  "//caffe2/torch/lib/c10d:c10d",
249  "//caffe2/torch/lib/libshm:libshm",
250  ],
251  external_deps=[
252  ("numpy", None, "cpp"),
253  ("pybind11", None),
254  ],
255  **common_flags
256  )
257 
258  return r