1 #pragma once 2 3 #include <ATen/core/Generator.h> 4 #include <ATen/detail/AcceleratorHooksInterface.h> 5 #include <c10/core/Allocator.h> 6 #include <c10/core/Device.h> 7 #include <c10/core/Storage.h> 8 #include <c10/util/Exception.h> 9 C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-parameter") 10 namespace at { 11 12 struct TORCH_API PrivateUse1HooksInterface : AcceleratorHooksInterface { 13 ~PrivateUse1HooksInterface() override = default; getDefaultGeneratorPrivateUse1HooksInterface14 virtual const at::Generator& getDefaultGenerator( 15 c10::DeviceIndex device_index) const { 16 TORCH_CHECK_NOT_IMPLEMENTED( 17 false, 18 "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `getDefaultGenerator`."); 19 } 20 getDeviceFromPtrPrivateUse1HooksInterface21 virtual at::Device getDeviceFromPtr(void* data) const { 22 TORCH_CHECK_NOT_IMPLEMENTED( 23 false, 24 "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `getDeviceFromPtr`."); 25 } 26 isPinnedPtrPrivateUse1HooksInterface27 virtual bool isPinnedPtr(const void* data) const override { 28 return false; 29 } 30 getPinnedMemoryAllocatorPrivateUse1HooksInterface31 virtual Allocator* getPinnedMemoryAllocator() const override { 32 TORCH_CHECK( 33 false, 34 "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `getPinnedMemoryAllocator`."); 35 } 36 hasPrimaryContextPrivateUse1HooksInterface37 virtual bool hasPrimaryContext(DeviceIndex device_index) const override { 38 TORCH_CHECK_NOT_IMPLEMENTED( 39 false, 40 "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `hasPrimaryContext`."); 41 } 42 initPrivateUse1PrivateUse1HooksInterface43 virtual void initPrivateUse1() const {} resizePrivateUse1BytesPrivateUse1HooksInterface44 virtual void resizePrivateUse1Bytes( 45 const c10::Storage& storage, 46 size_t newsize) const { 47 TORCH_CHECK_NOT_IMPLEMENTED( 48 false, 49 "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `resizePrivateUse1Bytes`."); 50 } 51 }; 52 53 struct TORCH_API PrivateUse1HooksArgs {}; 54 55 TORCH_API void RegisterPrivateUse1HooksInterface( 56 at::PrivateUse1HooksInterface* hook_); 57 58 TORCH_API bool isPrivateUse1HooksRegistered(); 59 60 namespace detail { 61 62 TORCH_API const at::PrivateUse1HooksInterface& getPrivateUse1Hooks(); 63 64 } // namespace detail 65 66 } // namespace at 67 C10_DIAGNOSTIC_POP() 68