libshm.h 1.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <ATen/MapAllocator.h>
  4. #ifdef __cplusplus
  5. #ifdef SHM_EXPORTS
  6. #define SHM_API __declspec(dllexport)
  7. #else
  8. #define SHM_API __declspec(dllimport)
  9. #endif
  10. SHM_API void libshm_init(const char* manager_exec_path);
  11. class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator {
  12. public:
  13. THManagedMapAllocator(
  14. const char* manager_handle,
  15. const char* filename,
  16. int flags,
  17. size_t size)
  18. : at::RefcountedMapAllocator(filename, flags, size) {}
  19. static at::DataPtr makeDataPtr(
  20. const char* manager_handle,
  21. const char* filename,
  22. int flags,
  23. size_t size);
  24. static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
  25. const char* manager_handle() const {
  26. return "no_manager";
  27. }
  28. };
  29. #endif
  30. #else
  31. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  32. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)