68 lines
1.9 KiB
C
68 lines
1.9 KiB
C
|
#pragma once
|
||
|
|
|
||
|
|
#include <atomic>
|
||
|
|
#include <mutex>
|
||
|
|
#include <utility>
|
||
|
|
|
||
|
|
#include <c10/macros/Macros.h>
|
||
|
|
#include <c10/util/C++17.h>
|
||
|
|
|
||
|
|
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:
|
||
|
|
#ifndef _WIN32
|
||
|
|
// 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
|
||
|
|
#endif
|
||
|
|
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
|