IPUHooksInterface.h 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. #pragma once
  2. #include <ATen/detail/AcceleratorHooksInterface.h>
  3. #include <c10/core/Allocator.h>
  4. #include <c10/util/Exception.h>
  5. #include <c10/util/Registry.h>
  6. namespace at {
  7. struct TORCH_API IPUHooksInterface : AcceleratorHooksInterface {
  8. ~IPUHooksInterface() override = default;
  9. void init() const override {
  10. TORCH_CHECK(false, "Cannot initialize IPU without ATen_ipu library.");
  11. }
  12. bool hasPrimaryContext(DeviceIndex /*device_index*/) const override {
  13. TORCH_CHECK(false, "Cannot initialize IPU without ATen_ipu library.");
  14. return false;
  15. }
  16. const Generator& getDefaultGenerator(
  17. [[maybe_unused]] DeviceIndex device_index = -1) const override {
  18. TORCH_CHECK(false, "Cannot initialize IPU without ATen_ipu library.");
  19. }
  20. Generator getNewGenerator(
  21. DeviceIndex /*device_index*/ = -1) const override {
  22. TORCH_CHECK(false, "Cannot initialize IPU without ATen_ipu library.");
  23. }
  24. };
  25. struct TORCH_API IPUHooksArgs {};
  26. TORCH_DECLARE_REGISTRY(IPUHooksRegistry, IPUHooksInterface, IPUHooksArgs);
  27. #define REGISTER_IPU_HOOKS(clsname) \
  28. C10_REGISTER_CLASS(IPUHooksRegistry, clsname, clsname)
  29. namespace detail {
  30. TORCH_API const IPUHooksInterface& getIPUHooks();
  31. } // namespace detail
  32. } // namespace at