IPUHooksInterface.h 1.5 KB

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