fallback registration if HPU not loaded

tbohutyn · bsochack · commit 1195682e82d8 · 2024-10-09T11:12:31.000+02:00
diff --git a/aten/src/ATen/detail/HPUHooksInterface.cpp b/aten/src/ATen/detail/HPUHooksInterface.cpp
@@ -1,19 +1,18 @@
 #include <ATen/detail/HPUHooksInterface.h>
 #include <c10/util/CallOnce.h>
+#include <memory>
 
 namespace at {
 namespace detail {
 
-static HPUHooksInterface* hpu_hooks = nullptr;
-
 TORCH_API const at::HPUHooksInterface& getHPUHooks() {
+  static std::unique_ptr<HPUHooksInterface> hpu_hooks;
   static c10::once_flag once;
   c10::call_once(once, [] {
-    hpu_hooks =
-        HPUHooksRegistry()->Create("HPUHooks", HPUHooksArgs{}).release();
-    TORCH_CHECK(
-        hpu_hooks != nullptr,
-        "Missing registration for HPUHooks in HPUHooksRegistry.")
+    hpu_hooks = HPUHooksRegistry()->Create("HPUHooks", HPUHooksArgs{});
+    if (!hpu_hooks) {
+      hpu_hooks = std::make_unique<HPUHooksInterface>();
+    }
   });
   return *hpu_hooks;
 }
diff --git a/aten/src/ATen/detail/HPUHooksInterface.h b/aten/src/ATen/detail/HPUHooksInterface.h
@@ -37,6 +37,11 @@ struct TORCH_API HPUHooksInterface : AcceleratorHooksInterface {
         false,
         "You should register `HPUHooksInterface` for HPU before call `getPinnedMemoryAllocator`.");
   }
+  bool hasPrimaryContext(C10_UNUSED DeviceIndex device_index) const override {
+    TORCH_CHECK(
+        false,
+        "You should register `HPUHooksInterface` for HPU before call `hasPrimaryContext`.");
+  }
 };
 
 struct TORCH_API HPUHooksArgs {};
