libshm.h 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. #pragma once
  2. #include <ATen/MapAllocator.h>
  3. #ifdef __cplusplus
  4. void libshm_init(const char* manager_exec_path);
  5. // Superclass to run a constructor before at::RefcountedMapAllocator
  6. class THManagedMapAllocatorInit {
  7. protected:
  8. THManagedMapAllocatorInit(const char* manager_handle, const char* filename);
  9. std::string manager_handle_;
  10. };
  11. // Like a at::RefcountedMapAllocator, but it also makes use of an external
  12. // shared memory manager process to ensure that shared memory regions actually
  13. // get freed in the end (even if processes lose the memory).
  14. class THManagedMapAllocator : private THManagedMapAllocatorInit,
  15. public at::RefcountedMapAllocator {
  16. public:
  17. THManagedMapAllocator(
  18. const char* manager_handle,
  19. const char* filename,
  20. int flags,
  21. ptrdiff_t size);
  22. void close() override;
  23. ~THManagedMapAllocator() override {
  24. close();
  25. }
  26. static at::DataPtr makeDataPtr(
  27. const char* manager_handle,
  28. const char* filename,
  29. int flags,
  30. ptrdiff_t size);
  31. static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
  32. const char* manager_handle() const {
  33. return manager_handle_.c_str();
  34. }
  35. };
  36. #endif