1 | #pragma once |
2 | |
3 | #include <atomic> |
4 | #include <mutex> |
5 | #include <thread> |
6 | #include <utility> |
7 | |
8 | #include <c10/macros/Macros.h> |
9 | #include <c10/util/C++17.h> |
10 | |
11 | namespace c10 { |
12 | |
13 | // custom c10 call_once implementation to avoid the deadlock in std::call_once. |
14 | // The implementation here is a simplified version from folly and likely much |
15 | // much higher memory footprint. |
16 | template <typename Flag, typename F, typename... Args> |
17 | inline void call_once(Flag& flag, F&& f, Args&&... args) { |
18 | if (C10_LIKELY(flag.test_once())) { |
19 | return; |
20 | } |
21 | flag.call_once_slow(std::forward<F>(f), std::forward<Args>(args)...); |
22 | } |
23 | |
24 | class once_flag { |
25 | public: |
26 | #ifndef _WIN32 |
27 | // running into build error on MSVC. Can't seem to get a repro locally so I'm |
28 | // just avoiding constexpr |
29 | // |
30 | // C:/actions-runner/_work/pytorch/pytorch\c10/util/CallOnce.h(26): error: |
31 | // defaulted default constructor cannot be constexpr because the |
32 | // corresponding implicitly declared default constructor would not be |
33 | // constexpr 1 error detected in the compilation of |
34 | // "C:/actions-runner/_work/pytorch/pytorch/aten/src/ATen/cuda/cub.cu". |
35 | constexpr |
36 | #endif |
37 | once_flag() noexcept = default; |
38 | once_flag(const once_flag&) = delete; |
39 | once_flag& operator=(const once_flag&) = delete; |
40 | |
41 | private: |
42 | template <typename Flag, typename F, typename... Args> |
43 | friend void call_once(Flag& flag, F&& f, Args&&... args); |
44 | |
45 | template <typename F, typename... Args> |
46 | void call_once_slow(F&& f, Args&&... args) { |
47 | std::lock_guard<std::mutex> guard(mutex_); |
48 | if (init_.load(std::memory_order_relaxed)) { |
49 | return; |
50 | } |
51 | c10::guts::invoke(f, std::forward<Args>(args)...); |
52 | init_.store(true, std::memory_order_release); |
53 | } |
54 | |
55 | bool test_once() { |
56 | return init_.load(std::memory_order_acquire); |
57 | } |
58 | |
59 | void reset_once() { |
60 | init_.store(false, std::memory_order_release); |
61 | } |
62 | |
63 | private: |
64 | std::mutex mutex_; |
65 | std::atomic<bool> init_{false}; |
66 | }; |
67 | |
68 | } // namespace c10 |
69 | |