HIPHooksInterface.h 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. #pragma once
  2. #include <c10/core/Allocator.h>
  3. #include <ATen/core/Generator.h>
  4. #include <c10/util/Exception.h>
  5. #include <c10/util/Registry.h>
  6. #include <cstddef>
  7. #include <functional>
  8. #include <memory>
  9. namespace at {
  10. class Context;
  11. }
  12. // NB: Class must live in `at` due to limitations of Registry.h.
  13. namespace at {
  14. // The HIPHooksInterface is an omnibus interface for any HIP functionality
  15. // which we may want to call into from CPU code (and thus must be dynamically
  16. // dispatched, to allow for separate compilation of HIP code). See
  17. // CUDAHooksInterface for more detailed motivation.
  18. struct TORCH_API HIPHooksInterface {
  19. // This should never actually be implemented, but it is used to
  20. // squelch -Werror=non-virtual-dtor
  21. virtual ~HIPHooksInterface() = default;
  22. // Initialize the HIP library state
  23. virtual void initHIP() const {
  24. AT_ERROR("Cannot initialize HIP without ATen_hip library.");
  25. }
  26. virtual std::unique_ptr<c10::GeneratorImpl> initHIPGenerator(Context*) const {
  27. AT_ERROR("Cannot initialize HIP generator without ATen_hip library.");
  28. }
  29. virtual bool hasHIP() const {
  30. return false;
  31. }
  32. virtual int64_t current_device() const {
  33. return -1;
  34. }
  35. virtual Allocator* getPinnedMemoryAllocator() const {
  36. AT_ERROR("Pinned memory requires HIP.");
  37. }
  38. virtual void registerHIPTypes(Context*) const {
  39. AT_ERROR("Cannot registerHIPTypes() without ATen_hip library.");
  40. }
  41. virtual int getNumGPUs() const {
  42. return 0;
  43. }
  44. };
  45. // NB: dummy argument to suppress "ISO C++11 requires at least one argument
  46. // for the "..." in a variadic macro"
  47. struct TORCH_API HIPHooksArgs {};
  48. C10_DECLARE_REGISTRY(HIPHooksRegistry, HIPHooksInterface, HIPHooksArgs);
  49. #define REGISTER_HIP_HOOKS(clsname) \
  50. C10_REGISTER_CLASS(HIPHooksRegistry, clsname, clsname)
  51. namespace detail {
  52. TORCH_API const HIPHooksInterface& getHIPHooks();
  53. } // namespace detail
  54. } // namespace at