1#pragma once
2
3#include <ATen/MapAllocator.h>
4
5#ifdef __cplusplus
6
7void libshm_init(const char* manager_exec_path);
8
9// Superclass to run a constructor before at::RefcountedMapAllocator
10class 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).
19class 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