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 namespace py = pybind11; 12 13 extern "C" { 14 15 #endif 16 17 // Flag to just run a frame normally 18 #define SKIP_CODE ((void*)0x1) 19 // Flag to run a frame and any recursive calls normally 20 #define SKIP_CODE_RECURSIVE ((void*)0x2) 21 22 // Points to the extra scratch space on the code object 23 extern Py_ssize_t extra_index; 24 25 // function to call when cache lookup errors 26 extern PyObject* guard_error_hook; 27 28 typedef PyObject FrameState; 29 typedef struct CacheEntry CacheEntry; 30 31 // ExtraState encasulates CacheEntry and FrameState. ExtraState is the highest 32 // level of abstraction of what is stored on the extra code object. Previously, 33 // we saved different parts on different extra indexes. We prefer this way 34 // because of cleaner abstraction and faster SetExtra access. 35 36 #ifdef __cplusplus 37 38 typedef struct VISIBILITY_HIDDEN ExtraState { 39 // List of cache entries for compiled code objects 40 std::list<CacheEntry> cache_entry_list; 41 // Frame state to detect dynamic shape dims 42 py::dict frame_state; 43 44 CacheEntry* get_first_entry(); 45 void move_to_front(CacheEntry* cache_entry); 46 void invalidate(CacheEntry* cache_entry); 47 } ExtraState; 48 49 #else 50 51 typedef struct ExtraState ExtraState; 52 53 #endif 54 55 // Helper to extra the cache_entry from the extra state. 56 // Ownership contract 57 // args 58 // - extra_state: Borrowed 59 // return 60 // - CacheEntry: Borrowed. 61 CacheEntry* extract_cache_entry(ExtraState* extra_state); 62 63 // Returns either the previously stored frame state or an empty dict. 64 // Ownership contract 65 // args 66 // - extra_state: Borrowed 67 // return 68 // - extra_state->frame_state: Borrowed. 69 FrameState* extract_frame_state(ExtraState* extra_state); 70 71 // Ownership contract 72 // args 73 // - code: Borrowed 74 // return 75 // - extra_state: Borrowed. 76 ExtraState* get_extra_state(PyCodeObject* code); 77 78 // This is passed as freefunc to _PyEval_RequestCodeExtraIndex. This acts as a 79 // deleter for the object on extra scratch space. This function is called 80 // internally in _PyCode_SetExtra and also during the code deallocation. 81 82 // Destroys the extra state by deleting cache_entry, frame state and finally 83 // freeing the constructed extra state. 84 85 // Developer note - You should not call this function directly. This is called 86 // directly inside set_extra_state. If you are in a situation trying to call 87 // this function, consider if set_extra_state should be called. 88 void destroy_extra_state(void* obj); 89 90 // Clears the existing object sitting on the extra scratch spance and sets it 91 // up with the new state. Note that _PyCode_SetExtra calls the 92 // destroy_extra_state deleter internally, and therefore we don't call it 93 // explicity here. 94 95 // Ownership contract 96 // args 97 // - extra_state: Stolen 98 // return 99 // - there is no return, but the extra_state is stolen, so it becomes 100 // set_extra_state responsibility to clean it up. It will be deleted during 101 // the reset_code/skip, when the set_extra_state is called with 102 // NULL/SKIP_CODE/SKIP_CODE_RECURSIVE. 103 104 // Invariant - Dont set the extra state for the extra state that is already on 105 // the code object. Otherwise, we will first free up the old extra state 106 // (which is also the new extra state) and write something invalid on the 107 // scratch space. 108 void set_extra_state(PyCodeObject* code, ExtraState* extra_state); 109 110 // Creates a new extra state and put it on the extra scrach space of the code 111 // object. 112 113 // Ownership contract 114 // args 115 // - code: Borrowed 116 // return: 117 // - extra_state: New reference. 118 // These references are then further passed to set_extra_state which becomes 119 // the final owner of these references. 120 ExtraState* init_and_set_extra_state(PyCodeObject* code); 121 122 // Lookup the cache held by extra_state. 123 // Ownership contract 124 // args 125 // - extra_state: Borrowed 126 // - f_locals: Borrowed 127 // return: 128 // - Py_None or PyCodeObject: Borrowed reference. 129 PyObject* lookup( 130 ExtraState* extra_state, 131 PyObject* f_locals, 132 PyObject* backend); 133 134 // Create a new cache entry at extra_state holding on to guarded_code. 135 // Ownership contract 136 // args 137 // - extra_state: Borrowed 138 // - guarded_code: Borrowed 139 // return: 140 // - cache_entry: Borrowed reference 141 CacheEntry* create_cache_entry( 142 ExtraState* extra_state, 143 PyObject* guraded_code, 144 PyObject* callback); 145 146 // Extracts the backend fn from the callback. 147 PyObject* get_backend(PyObject* callback); 148 149 #ifdef __cplusplus 150 151 } // extern "C" 152 153 // Returns the list of CacheEntry corresponding to code_obj. 154 // Warning: returns references whose lifetimes are controlled by C++ 155 py::list _debug_get_cache_entry_list(const py::handle& code_obj); 156 157 #endif 158