libshm.h 779 B

123456789101112131415161718192021222324252627282930313233343536
  1. #pragma once
  2. #include <ATen/MapAllocator.h>
  3. #ifdef __cplusplus
  4. #ifdef SHM_EXPORTS
  5. #define SHM_API __declspec(dllexport)
  6. #else
  7. #define SHM_API __declspec(dllimport)
  8. #endif
  9. SHM_API void libshm_init(const char* manager_exec_path);
  10. class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator {
  11. public:
  12. THManagedMapAllocator(
  13. const char* manager_handle,
  14. const char* filename,
  15. int flags,
  16. size_t size)
  17. : at::RefcountedMapAllocator(filename, flags, size) {}
  18. static at::DataPtr makeDataPtr(
  19. const char* manager_handle,
  20. const char* filename,
  21. int flags,
  22. size_t size);
  23. static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
  24. const char* manager_handle() const {
  25. return "no_manager";
  26. }
  27. };
  28. #endif