PrivateUse1HooksInterface.h 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. #pragma once
  2. #include <ATen/core/GeneratorForPrivateuseone.h>
  3. #include <ATen/detail/AcceleratorHooksInterface.h>
  4. #include <c10/core/Allocator.h>
  5. #include <c10/core/Device.h>
  6. #include <c10/core/Storage.h>
  7. #include <c10/util/Exception.h>
  8. C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-parameter")
  9. namespace at {
  10. struct TORCH_API PrivateUse1HooksInterface : AcceleratorHooksInterface {
  11. #define FAIL_PRIVATEUSE1HOOKS_FUNC(func) \
  12. TORCH_CHECK_NOT_IMPLEMENTED( \
  13. false, \
  14. "You should register `PrivateUse1HooksInterface`", \
  15. "by `RegisterPrivateUse1HooksInterface` and implement `", \
  16. func, \
  17. "` at the same time for PrivateUse1.");
  18. ~PrivateUse1HooksInterface() override = default;
  19. bool isBuilt() const override {
  20. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  21. }
  22. bool isAvailable() const override {
  23. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  24. }
  25. const at::Generator& getDefaultGenerator(
  26. c10::DeviceIndex device_index) const override {
  27. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  28. }
  29. Generator getNewGenerator(
  30. [[maybe_unused]] DeviceIndex device_index = -1) const override {
  31. // TODO(FFFrog): Perserved for BC and will be removed in the future.
  32. if (at::GetGeneratorPrivate().has_value())
  33. return at::GetGeneratorForPrivateuse1(device_index);
  34. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  35. }
  36. at::Device getDeviceFromPtr(void* data) const override {
  37. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  38. }
  39. bool isPinnedPtr(const void* data) const override {
  40. return false;
  41. }
  42. Allocator* getPinnedMemoryAllocator() const override {
  43. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  44. }
  45. bool hasPrimaryContext(DeviceIndex device_index) const override {
  46. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  47. }
  48. void init() const override {}
  49. virtual void resizePrivateUse1Bytes(
  50. const c10::Storage& storage,
  51. size_t newsize) const {
  52. FAIL_PRIVATEUSE1HOOKS_FUNC(__func__);
  53. }
  54. #undef FAIL_PRIVATEUSE1HOOKS_FUNC
  55. };
  56. struct TORCH_API PrivateUse1HooksArgs {};
  57. TORCH_API void RegisterPrivateUse1HooksInterface(
  58. at::PrivateUse1HooksInterface* hook_);
  59. TORCH_API bool isPrivateUse1HooksRegistered();
  60. namespace detail {
  61. TORCH_API const at::PrivateUse1HooksInterface& getPrivateUse1Hooks();
  62. } // namespace detail
  63. } // namespace at
  64. C10_DIAGNOSTIC_POP()