1 #pragma once 2 3 #include <c10/macros/Export.h> 4 #include <atomic> 5 6 namespace c10::impl { 7 8 // This TLS controls whether or not we permanently associate PyObject 9 // with Tensor the first time it is allocated. When hermetic PyObject 10 // TLS is enabled (state is true), we DO NOT save PyObjects to Tensor, 11 // meaning you get a distinct PyObject whenever you execute the code in 12 // question. 13 struct C10_API HermeticPyObjectTLS { 14 static void set_state(bool state); get_stateHermeticPyObjectTLS15 static bool get_state() { 16 // Hypothetical fastpath if torchdeploy/multipy isn't used. Per 17 // https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2020/p2055r0.pdf 18 // this qualifies relaxed access because it is a single-location data 19 // structure (only the boolean here). 20 // 21 // Forgetting about data races for a moment, is there a logical race? 22 // 23 // - Boolean only ever transitions from false to true. So the 24 // critical situation is when one interpreter is already running 25 // when a second interpreter switches haveState from false to true. 26 // 27 // - The first interpreter is indifferent whether or not it sees 28 // hasState true/false; obviously false works (this is what the 29 // interpreter was previously using; more directly, the interpreter 30 // calls into itself as the handler, so being hermetic is not 31 // required), and true simply means serviced python operator calls will 32 // be hermetic; in these cases it is expected to be functionally 33 // equivalent. 34 // 35 // - The second interpreter MUST see hasState true (as its requests will 36 // be forwarded to the first interpreter), but it is assumed that there 37 // is a synchronization between the interpreter initialization, and 38 // when we actually perform operations, so it is guaranteed to see 39 // hasState true. 40 // 41 // QED. 42 // 43 // This fastpath is currently disabled so that we can more easily test that 44 // hermetic mode works correctly even on stock build of PyTorch. 45 if (false && !haveState_.load(std::memory_order_relaxed)) 46 return false; 47 return get_tls_state(); 48 } 49 // Call this from the multipy/torchdeploy top level 50 static void init_state(); 51 52 private: 53 // This only flipped once from false to true during torchdeploy/multipy 54 // initialization, and never again. 55 static std::atomic<bool> haveState_; 56 static bool get_tls_state(); 57 }; 58 59 } // namespace c10::impl 60