1 | #pragma once |
2 | |
3 | #include <c10/core/Allocator.h> |
4 | |
5 | namespace at { |
6 | |
7 | enum MappedAllocatorModes { |
8 | ALLOCATOR_MAPPED_SHARED = 1, |
9 | ALLOCATOR_MAPPED_SHAREDMEM = 2, |
10 | ALLOCATOR_MAPPED_EXCLUSIVE = 4, |
11 | ALLOCATOR_MAPPED_NOCREATE = 8, |
12 | ALLOCATOR_MAPPED_KEEPFD = 16, |
13 | ALLOCATOR_MAPPED_FROMFD = 32, |
14 | ALLOCATOR_MAPPED_UNLINK = 64 |
15 | }; |
16 | |
17 | // Sentinel value/type to help distinguish the file descriptor constructor from |
18 | // the non-file descriptor constructor |
19 | enum WithFd { WITH_FD }; |
20 | |
21 | TORCH_API std::string NewProcessWideShmHandle(); |
22 | |
23 | class TORCH_API MapAllocator { |
24 | public: |
25 | MapAllocator(std::string filename, int flags, size_t size); |
26 | MapAllocator(WithFd, std::string filename, int fd, int flags, size_t size); |
27 | MapAllocator(const MapAllocator&) = delete; |
28 | MapAllocator& operator=(const MapAllocator&) = delete; |
29 | MapAllocator(MapAllocator&&) = delete; |
30 | MapAllocator& operator=(MapAllocator&&) = delete; |
31 | |
32 | const char* filename() const { |
33 | return filename_.c_str(); |
34 | } |
35 | int fd() const { |
36 | #ifdef _WIN32 |
37 | TORCH_CHECK(false, "MapAllocator::fd() is unsupported on Windows" ); |
38 | #else |
39 | return fd_; |
40 | #endif |
41 | } |
42 | ptrdiff_t size() const { |
43 | return size_; |
44 | } |
45 | // Return a pointer to the actual data for this allocator |
46 | // (in the case of the refcounted allocator, this is offset |
47 | // from the base pointer.) |
48 | virtual void* data() const { |
49 | return base_ptr_; |
50 | } |
51 | |
52 | static MapAllocator* fromDataPtr(const at::DataPtr&); |
53 | static at::DataPtr makeDataPtr( |
54 | std::string filename, |
55 | int flags, |
56 | size_t size, |
57 | size_t* actual_size_out); |
58 | static at::DataPtr makeDataPtr( |
59 | WithFd, |
60 | const char* filename, |
61 | int fd, |
62 | int flags, |
63 | size_t size, |
64 | size_t* actual_size_out); |
65 | |
66 | // Closes the data. Helps us avoid destructor shenanigans |
67 | virtual void close(); |
68 | |
69 | // This is very dangerous. You have to redefine this destructor for each |
70 | // subclass |
71 | virtual ~MapAllocator(); |
72 | |
73 | protected: |
74 | bool closed_ = false; |
75 | std::string filename_; |
76 | int flags_ = 0; |
77 | ptrdiff_t size_; /* mapped size */ |
78 | #ifdef _WIN32 |
79 | void* handle_; |
80 | void* event_; |
81 | std::string eventname_; |
82 | #else |
83 | int fd_ = -1; |
84 | #endif |
85 | void* base_ptr_ = nullptr; |
86 | }; |
87 | |
88 | // Base-from-member idiom |
89 | struct TORCH_API RefcountedMapAllocatorArgCheck { |
90 | RefcountedMapAllocatorArgCheck(int flags); |
91 | }; |
92 | |
93 | class TORCH_API RefcountedMapAllocator : private RefcountedMapAllocatorArgCheck, |
94 | public MapAllocator { |
95 | public: |
96 | RefcountedMapAllocator(const char* filename, int flags, size_t size); |
97 | RefcountedMapAllocator( |
98 | WithFd, |
99 | const char* filename, |
100 | int fd, |
101 | int flags, |
102 | size_t size); |
103 | |
104 | static RefcountedMapAllocator* fromDataPtr(const at::DataPtr&); |
105 | static at::DataPtr makeDataPtr( |
106 | const char* filename, |
107 | int flags, |
108 | size_t size, |
109 | size_t* actual_size_out); |
110 | static at::DataPtr makeDataPtr( |
111 | WithFd, |
112 | const char* filename, |
113 | int fd, |
114 | int flags, |
115 | size_t size, |
116 | size_t* actual_size_out); |
117 | |
118 | void* data() const override; |
119 | |
120 | void incref(); |
121 | int decref(); |
122 | void close() override; |
123 | |
124 | ~RefcountedMapAllocator() override { |
125 | close(); |
126 | } |
127 | |
128 | protected: |
129 | void checkFlags(); |
130 | void initializeAlloc(); |
131 | }; |
132 | |
133 | } // namespace at |
134 | |