1 | #include <torch/csrc/utils/cuda_lazy_init.h> |
---|---|
2 | |
3 | #include <torch/csrc/Exceptions.h> |
4 | #include <torch/csrc/python_headers.h> |
5 | #include <torch/csrc/utils/object_ptr.h> |
6 | |
7 | namespace torch { |
8 | namespace utils { |
9 | namespace { |
10 | |
11 | bool is_initialized = false; |
12 | |
13 | } |
14 | |
15 | void cuda_lazy_init() { |
16 | pybind11::gil_scoped_acquire g; |
17 | // Protected by the GIL. We don't use call_once because under ASAN it |
18 | // has a buggy implementation that deadlocks if an instance throws an |
19 | // exception. In any case, call_once isn't necessary, because we |
20 | // have taken a lock. |
21 | if (is_initialized) { |
22 | return; |
23 | } |
24 | |
25 | auto module = THPObjectPtr(PyImport_ImportModule("torch.cuda")); |
26 | if (!module) { |
27 | throw python_error(); |
28 | } |
29 | |
30 | auto res = THPObjectPtr(PyObject_CallMethod(module.get(), "_lazy_init", "")); |
31 | if (!res) { |
32 | throw python_error(); |
33 | } |
34 | |
35 | is_initialized = true; |
36 | } |
37 | |
38 | void set_requires_cuda_init(bool value) { |
39 | is_initialized = !value; |
40 | } |
41 | |
42 | } // namespace utils |
43 | } // namespace torch |
44 |