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