1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
|
#include <torch/csrc/dynamo/init.h>
#include <torch/csrc/dynamo/utils.h>
#include <pybind11/stl_bind.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/dynamo/cache_entry.h>
#include <torch/csrc/dynamo/cpython_defs.h>
#include <torch/csrc/dynamo/eval_frame.h>
#include <torch/csrc/dynamo/extra_state.h>
#include <torch/csrc/dynamo/guards.h>
#include <torch/csrc/dynamo/python_compiled_autograd.h>
#include <torch/csrc/utils/pybind.h>
#include <torch/csrc/utils/python_compat.h>
static struct PyModuleDef _module =
{PyModuleDef_HEAD_INIT, "torch._C._dynamo", "", -1, nullptr};
PYBIND11_MAKE_OPAQUE(std::vector<uint8_t>)
namespace torch::dynamo {
#if IS_PYTHON_3_11_PLUS
std::vector<uint8_t> _PyOpcode_Caches_vec(
THP_PyOpcode_Caches,
THP_PyOpcode_Caches + THP_PyOpcode_Caches_size);
#else
std::vector<uint8_t> _PyOpcode_Caches_vec;
#endif
using torch::dynamo::autograd::torch_c_dynamo_compiled_autograd_init;
void initDynamoBindings(PyObject* torch) {
PyObject* dynamo = PyModule_Create(&_module);
if (dynamo == nullptr || PyModule_AddObject(torch, "_dynamo", dynamo) != 0) {
throw python_error();
}
#ifdef Py_GIL_DISABLED
PyUnstable_Module_SetGIL(dynamo, Py_MOD_GIL_NOT_USED);
#endif
PyObject* eval_frame = torch_c_dynamo_eval_frame_init();
if (eval_frame == nullptr ||
PyModule_AddObject(dynamo, "eval_frame", eval_frame) != 0) {
throw python_error();
}
PyObject* utils = torch_c_dynamo_utils_init();
if (utils == nullptr || PyModule_AddObject(dynamo, "utils", utils) != 0) {
throw python_error();
}
PyObject* guards = torch_c_dynamo_guards_init();
if (guards == nullptr || PyModule_AddObject(dynamo, "guards", guards) != 0) {
throw python_error();
}
PyObject* compiled_autograd = torch_c_dynamo_compiled_autograd_init();
if (compiled_autograd == nullptr ||
PyModule_AddObject(dynamo, "compiled_autograd", compiled_autograd) != 0) {
throw python_error();
}
auto m = py::handle(eval_frame).cast<py::module>();
py::class_<CacheEntry>(m, "_CacheEntry")
.def_readonly("guard_manager", &CacheEntry::guard_manager)
.def_readonly("code", &CacheEntry::code)
.def_readonly("compile_id", &CacheEntry::compile_id)
.def_readonly("trace_annotation", &CacheEntry::trace_annotation)
.def_property_readonly("next", &CacheEntry::next)
.def(
"update_diff_guard_root_manager",
&CacheEntry::update_diff_guard_root_manager);
py::class_<ExtraState>(m, "_ExtraState")
.def("invalidate", &ExtraState::invalidate);
m.def("_debug_get_cache_entry_list", &_debug_get_cache_entry_list);
py::bind_vector<std::vector<uint8_t>>(m, "VectorUInt8");
m.attr("py_opcode_caches") = _PyOpcode_Caches_vec;
}
} // namespace torch::dynamo
|