1 #pragma once 2 3 #include <Python.h> 4 5 #ifdef __cplusplus 6 7 #include <torch/csrc/dynamo/utils.h> 8 #include <torch/csrc/utils/pybind.h> 9 #include <list> 10 11 extern "C" { 12 13 #endif 14 15 /* 16 Our cache resides on the extra scratch space of the code object. The structure 17 of the cache is as follows: 18 19 -> ExtraState 20 -> CacheEntry (list) 21 -> check_fn 22 -> code 23 -> FrameState 24 25 CacheEntry is a linked list node containing the check_fn for guards 26 and the optimized code. 27 28 The FrameState is a PyDict that enables sharing between different frames. This 29 is used to detect dynamism in automatic dynamic shapes. 30 31 These two are encapsulated into a ExtraState. 32 */ 33 34 typedef struct CacheEntry CacheEntry; 35 typedef struct ExtraState ExtraState; 36 37 #ifdef __cplusplus 38 39 typedef struct VISIBILITY_HIDDEN CacheEntry { 40 // check the guards: lambda: <locals of user function>: bool 41 py::object check_fn; 42 // modified user bytecode (protected by check_fn's guards) 43 py::object code; 44 // CompileId corresponding to this compilation 45 py::object compile_id; 46 // root guard manager if exists 47 void* root_mgr{nullptr}; 48 // backend used to create this cache entry 49 PyObject* backend{nullptr}; 50 // Reference to owning ExtraState 51 ExtraState* _owner{nullptr}; 52 // Reference to this CacheEntry's location in owner's linked list 53 std::list<CacheEntry>::iterator _owner_loc; 54 55 CacheEntry(const py::handle& guarded_code, PyObject* backend); 56 ~CacheEntry(); 57 58 // Warning: returns a reference whose lifetime is controlled by C++ 59 py::object next(); 60 } CacheEntry; 61 62 #endif 63 64 // Returns borrowed reference 65 PyCodeObject* CacheEntry_get_code(CacheEntry* e); 66 67 // Returns a borrowed reference to CacheEntry as a PyObject 68 // Warning: lifetime is controlled by C++ 69 PyObject* CacheEntry_to_obj(CacheEntry* e); 70 71 #ifdef __cplusplus 72 } // extern "C" 73 #endif 74