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 | ptrdiff_t size); |
27 | |
28 | void close() override; |
29 | |
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 | ptrdiff_t size); |
39 | static THManagedMapAllocator* fromDataPtr(const at::DataPtr&); |
40 | |
41 | const char* manager_handle() const { |
42 | return manager_handle_.c_str(); |
43 | } |
44 | }; |
45 | |
46 | #endif |
47 | |