diff --git a/torch/csrc/dynamo/cpython_defs.c b/torch/csrc/dynamo/cpython_defs.c index b40efc21ad0a..7b86017c59b3 100644 --- a/torch/csrc/dynamo/cpython_defs.c +++ b/torch/csrc/dynamo/cpython_defs.c @@ -5,13 +5,15 @@ #if IS_PYTHON_3_15_PLUS || (IS_PYTHON_3_14_PLUS && defined(_WIN32)) const uint8_t* THP_PyOpcode_Caches = NULL; -const int THP_PyOpcode_Caches_size = 0; +int THP_PyOpcode_Caches_size = 0; void THP_PyThreadState_PopFrame( PyThreadState* tstate, _PyInterpreterFrame* frame) {} void THP_PyFrame_Clear(_PyInterpreterFrame* frame) {} +void init_THPCaches() {} + #else #if IS_PYTHON_3_11_PLUS @@ -481,16 +483,13 @@ void THP_PyThreadState_PopFrame( #endif -#if IS_PYTHON_3_11_PLUS - -const uint8_t* THP_PyOpcode_Caches = _PyOpcode_Caches; -const int THP_PyOpcode_Caches_size = sizeof(_PyOpcode_Caches) / sizeof(uint8_t); - -#else - const uint8_t* THP_PyOpcode_Caches = NULL; -const int THP_PyOpcode_Caches_size = 0; - +int THP_PyOpcode_Caches_size = 0; +void init_THPCaches() { +#if IS_PYTHON_3_11_PLUS + THP_PyOpcode_Caches = _PyOpcode_Caches; + THP_PyOpcode_Caches_size = sizeof(_PyOpcode_Caches) / sizeof(uint8_t); #endif +} #endif // IS_PYTHON_3_15_PLUS diff --git a/torch/csrc/dynamo/cpython_defs.h b/torch/csrc/dynamo/cpython_defs.h index 5a58c7ee8c77..7183875dc682 100644 --- a/torch/csrc/dynamo/cpython_defs.h +++ b/torch/csrc/dynamo/cpython_defs.h @@ -28,13 +28,13 @@ void THP_PyThreadState_PopFrame( // pointers to _PyOpcode_Caches for C++ #ifdef __cplusplus - -extern "C" const uint8_t* THP_PyOpcode_Caches; -extern "C" const int THP_PyOpcode_Caches_size; - -#else +extern "C" { +#endif extern const uint8_t* THP_PyOpcode_Caches; -extern const int THP_PyOpcode_Caches_size; +extern int THP_PyOpcode_Caches_size; +void init_THPCaches(); +#ifdef __cplusplus +} // extern "C" #endif diff --git a/torch/csrc/dynamo/init.cpp b/torch/csrc/dynamo/init.cpp index 13826bf7f356..f1590e19d49c 100644 --- a/torch/csrc/dynamo/init.cpp +++ b/torch/csrc/dynamo/init.cpp @@ -21,18 +21,8 @@ PYBIND11_MAKE_OPAQUE(std::vector) namespace torch::dynamo { -#if IS_PYTHON_3_11_PLUS - -std::vector _PyOpcode_Caches_vec( - THP_PyOpcode_Caches, - THP_PyOpcode_Caches + THP_PyOpcode_Caches_size); - -#else - std::vector _PyOpcode_Caches_vec; -#endif - using torch::dynamo::autograd::torch_c_dynamo_compiled_autograd_init; namespace { @@ -265,6 +255,13 @@ void initDynamoBindings(PyObject* torch) { m.def("_load_precompile_entry", &_load_precompile_entry); m.def("_debug_get_precompile_entries", &_debug_get_precompile_entries); py::bind_vector>(m, "VectorUInt8"); + init_THPCaches(); + if (THP_PyOpcode_Caches != nullptr) { + _PyOpcode_Caches_vec.insert( + _PyOpcode_Caches_vec.end(), + THP_PyOpcode_Caches, + THP_PyOpcode_Caches + THP_PyOpcode_Caches_size); + } m.attr("py_opcode_caches") = _PyOpcode_Caches_vec; m.def("code_framelocals_names", &code_framelocals_names); _register_functions(dynamo);