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
11namespace 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.
16template <typename Flag, typename F, typename... Args>
17inline 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
24class 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