Spaces:
Running
Running
namespace c10 { | |
// custom c10 call_once implementation to avoid the deadlock in std::call_once. | |
// The implementation here is a simplified version from folly and likely much | |
// much higher memory footprint. | |
template <typename Flag, typename F, typename... Args> | |
inline void call_once(Flag& flag, F&& f, Args&&... args) { | |
if (C10_LIKELY(flag.test_once())) { | |
return; | |
} | |
flag.call_once_slow(std::forward<F>(f), std::forward<Args>(args)...); | |
} | |
class once_flag { | |
public: | |
// running into build error on MSVC. Can't seem to get a repro locally so I'm | |
// just avoiding constexpr | |
// | |
// C:/actions-runner/_work/pytorch/pytorch\c10/util/CallOnce.h(26): error: | |
// defaulted default constructor cannot be constexpr because the | |
// corresponding implicitly declared default constructor would not be | |
// constexpr 1 error detected in the compilation of | |
// "C:/actions-runner/_work/pytorch/pytorch/aten/src/ATen/cuda/cub.cu". | |
constexpr | |
once_flag() noexcept = default; | |
once_flag(const once_flag&) = delete; | |
once_flag& operator=(const once_flag&) = delete; | |
private: | |
template <typename Flag, typename F, typename... Args> | |
friend void call_once(Flag& flag, F&& f, Args&&... args); | |
template <typename F, typename... Args> | |
void call_once_slow(F&& f, Args&&... args) { | |
std::lock_guard<std::mutex> guard(mutex_); | |
if (init_.load(std::memory_order_relaxed)) { | |
return; | |
} | |
c10::guts::invoke(std::forward<F>(f), std::forward<Args>(args)...); | |
init_.store(true, std::memory_order_release); | |
} | |
bool test_once() { | |
return init_.load(std::memory_order_acquire); | |
} | |
void reset_once() { | |
init_.store(false, std::memory_order_release); | |
} | |
private: | |
std::mutex mutex_; | |
std::atomic<bool> init_{false}; | |
}; | |
} // namespace c10 | |