xref: /aosp_15_r20/external/pytorch/aten/src/ATen/detail/PrivateUse1HooksInterface.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <ATen/detail/PrivateUse1HooksInterface.h>
2 
3 namespace at {
4 
5 static PrivateUse1HooksInterface* privateuse1_hooks = nullptr;
6 static std::mutex _hooks_mutex_lock;
7 
RegisterPrivateUse1HooksInterface(at::PrivateUse1HooksInterface * hook_)8 TORCH_API void RegisterPrivateUse1HooksInterface(at::PrivateUse1HooksInterface* hook_) {
9   std::lock_guard<std::mutex> lock(_hooks_mutex_lock);
10   TORCH_CHECK(privateuse1_hooks == nullptr, "PrivateUse1HooksInterface only could be registered once.");
11   privateuse1_hooks = hook_;
12 }
13 
isPrivateUse1HooksRegistered()14 TORCH_API bool isPrivateUse1HooksRegistered() {
15   return privateuse1_hooks != nullptr;
16 }
17 
18 namespace detail {
19 
getPrivateUse1Hooks()20 TORCH_API const at::PrivateUse1HooksInterface& getPrivateUse1Hooks() {
21   TORCH_CHECK(
22       privateuse1_hooks != nullptr,
23       "Please register PrivateUse1HooksInterface by `RegisterPrivateUse1HooksInterface` first.");
24   return *privateuse1_hooks;
25 }
26 
27 } // namespace detail
28 
29 } // namespace at
30