1#pragma once
2
3#include <c10/core/Allocator.h>
4
5namespace at {
6
7enum 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
19enum WithFd { WITH_FD };
20
21TORCH_API std::string NewProcessWideShmHandle();
22
23class 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
89struct TORCH_API RefcountedMapAllocatorArgCheck {
90 RefcountedMapAllocatorArgCheck(int flags);
91};
92
93class 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