ORTHooksInterface.h 1.1 KB

123456789101112131415161718192021222324252627282930313233343536
  1. #pragma once
  2. #include <c10/util/Exception.h>
  3. #include <c10/util/Registry.h>
  4. constexpr const char* ORT_HELP =
  5. " You need to 'import torch_ort' to use the 'ort' device in PyTorch. "
  6. "The 'torch_ort' module is provided by the ONNX Runtime itself "
  7. "(https://onnxruntime.ai).";
  8. // NB: Class must live in `at` due to limitations of Registry.h.
  9. namespace at {
  10. struct TORCH_API ORTHooksInterface {
  11. // This should never actually be implemented, but it is used to
  12. // squelch -Werror=non-virtual-dtor
  13. virtual ~ORTHooksInterface() = default;
  14. virtual std::string showConfig() const {
  15. TORCH_CHECK(false, "Cannot query detailed ORT version information.", ORT_HELP);
  16. }
  17. };
  18. // NB: dummy argument to suppress "ISO C++11 requires at least one argument
  19. // for the "..." in a variadic macro"
  20. struct TORCH_API ORTHooksArgs {};
  21. C10_DECLARE_REGISTRY(ORTHooksRegistry, ORTHooksInterface, ORTHooksArgs);
  22. #define REGISTER_ORT_HOOKS(clsname) \
  23. C10_REGISTER_CLASS(ORTHooksRegistry, clsname, clsname)
  24. namespace detail {
  25. TORCH_API const ORTHooksInterface& getORTHooks();
  26. } // namespace detail
  27. } // namespace at