1 #pragma once 2 3 #include <ATen/core/Generator.h> 4 #include <c10/core/Allocator.h> 5 #include <c10/util/Exception.h> 6 #include <c10/util/Registry.h> 7 8 namespace at { 9 10 struct TORCH_API IPUHooksInterface { 11 virtual ~IPUHooksInterface() = default; 12 13 virtual const Generator& getDefaultIPUGenerator( 14 DeviceIndex device_index [[maybe_unused]] = -1) const { 15 AT_ERROR( 16 "Cannot get the default IPU generator: the IPU backend is not " 17 "available."); 18 } 19 20 virtual Generator newIPUGenerator(DeviceIndex device_index [[maybe_unused]] = -1) const { 21 AT_ERROR( 22 "Cannot create a new IPU generator: the IPU backend is not available."); 23 } 24 }; 25 26 struct TORCH_API IPUHooksArgs {}; 27 28 TORCH_DECLARE_REGISTRY(IPUHooksRegistry, IPUHooksInterface, IPUHooksArgs); 29 #define REGISTER_IPU_HOOKS(clsname) \ 30 C10_REGISTER_CLASS(IPUHooksRegistry, clsname, clsname) 31 32 namespace detail { 33 TORCH_API const IPUHooksInterface& getIPUHooks(); 34 } // namespace detail 35 } // namespace at 36