1 | // Copyright 2021 Google LLC |
2 | // |
3 | // This source code is licensed under the BSD-style license found in the |
4 | // LICENSE file in the root directory of this source tree. |
5 | |
6 | // Include first for the platform detection macros. |
7 | #include "xnnpack/common.h" |
8 | |
9 | #if XNN_PLATFORM_WINDOWS |
10 | #ifndef WIN32_LEAN_AND_MEAN |
11 | #define WIN32_LEAN_AND_MEAN |
12 | #endif |
13 | #include <windows.h> |
14 | #else |
15 | // This define needs to come first because errno include features.h and would have defined macros that lead to |
16 | // sys/mman.h not having mremap. |
17 | #if !defined(_GNU_SOURCE) |
18 | #define _GNU_SOURCE |
19 | #endif |
20 | #include <errno.h> |
21 | #include <sys/mman.h> |
22 | #include <unistd.h> |
23 | #endif |
24 | |
25 | #include <stddef.h> |
26 | #include <stdint.h> |
27 | #include <xnnpack.h> |
28 | |
29 | #include "xnnpack/allocator.h" |
30 | #include "xnnpack/log.h" |
31 | #include "xnnpack/math.h" |
32 | #include "xnnpack/params.h" |
33 | |
34 | // Helpers to allocate/mmap and release memory used by both code and weights cache. |
35 | |
36 | // Maps `size` bytes of memory, returns pointer to allocation, NULL if failed. |
37 | static void* allocate_buffer(size_t size) { |
38 | xnn_log_debug("allocating buffer of size %zu" , size); |
39 | assert(size == round_up_po2(size, xnn_params.page_size)); |
40 | #if XNN_PLATFORM_WINDOWS |
41 | void* p = VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); |
42 | if (p == NULL) { |
43 | xnn_log_error("failed to allocate %zu bytes for code/weights buffer, error code: %" PRIu32, |
44 | size, (uint32_t) GetLastError()); |
45 | return NULL; |
46 | } |
47 | #else |
48 | void* p = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); |
49 | if (p == MAP_FAILED) { |
50 | xnn_log_error("failed to allocate %zu bytes for code/weights buffer, error code: %d" , size, errno); |
51 | return NULL; |
52 | } |
53 | #endif |
54 | return p; |
55 | } |
56 | |
57 | // Releases memory previously mapped by `allocate_buffer`, returns xnn_status_success on success. |
58 | static enum xnn_status release_memory(void* start, size_t capacity) { |
59 | #if XNN_PLATFORM_WINDOWS |
60 | // We only decommited any unused capacity, so we release all of it now. |
61 | if (!VirtualFree(start, 0, MEM_RELEASE)) { |
62 | xnn_log_error("failed to release code/weights buffer, error code: %" PRIu32, (uint32_t) GetLastError()); |
63 | return xnn_status_invalid_state; |
64 | } |
65 | #else |
66 | if (munmap(start, capacity) == -1) { |
67 | xnn_log_error("failed to release code/weights buffer, error code: %d" , errno); |
68 | return xnn_status_invalid_state; |
69 | } |
70 | #endif |
71 | return xnn_status_success; |
72 | } |
73 | |
74 | // Resize a buffer at old_pointer of size old_bytes to new_size. The actual new size of the resized buffer is written to |
75 | // new_capacity_out, which can be >= new_size due to page alignment requirements. |
76 | // Returns a pointer to a buffer which might be the same as old_pointer if we can remap virtual memory, otherwise we |
77 | // allocate a new buffer and copy contents of old_buffer over. |
78 | static void* resize_buffer( |
79 | void* old_pointer, size_t old_size, size_t old_capacity, size_t new_size, size_t* new_capacity_out) |
80 | { |
81 | size_t new_capacity = round_up_po2(new_size, xnn_params.page_size); |
82 | #if XNN_PLATFORM_LINUX |
83 | void* new_pointer = mremap(old_pointer, old_size, new_capacity, MREMAP_MAYMOVE, NULL); |
84 | if (new_pointer == MAP_FAILED) { |
85 | xnn_log_error("mremap failed with errno: %d" , errno); |
86 | return NULL; |
87 | } |
88 | xnn_log_debug("resize_buffer: remap, old capacity %zu to new capacity %zu" , old_capacity, new_capacity); |
89 | #else |
90 | void* new_pointer = allocate_buffer(new_capacity); |
91 | if (new_pointer == NULL) { |
92 | xnn_log_error("allocate_buffer failed" ); |
93 | return NULL; |
94 | } |
95 | memcpy(new_pointer, old_pointer, old_size); |
96 | // Release old code_buffer. |
97 | enum xnn_status status = release_memory(old_pointer, old_capacity); |
98 | if (status != xnn_status_success) { |
99 | xnn_log_error("releasing old buffer failed, this could be a leak of %zu bytes" , old_capacity); |
100 | // Log but proceed as per normal since we successfully allocated a new memory that can be used by the caller. |
101 | } |
102 | xnn_log_debug("resize_buffer: allocate memory, old capacity %zu to new capacity %zu" , old_capacity, new_capacity); |
103 | #endif |
104 | *new_capacity_out = new_capacity; |
105 | return new_pointer; |
106 | } |
107 | |
108 | enum xnn_status xnn_allocate_code_memory(struct xnn_code_buffer* buf, size_t size) { |
109 | memset(buf, 0, sizeof(struct xnn_code_buffer)); |
110 | size_t page_aligned_size = round_up_po2(size, xnn_params.page_size); |
111 | buf->start = allocate_buffer(page_aligned_size); |
112 | if (buf->start == NULL) { |
113 | return xnn_status_out_of_memory; |
114 | } |
115 | |
116 | buf->size = 0; |
117 | buf->capacity = page_aligned_size; |
118 | return xnn_status_success; |
119 | } |
120 | |
121 | // Releases unused memory. Will write the new capacity to `capacity`. |
122 | static enum xnn_status release_unused_memory(size_t size, void* start, size_t* capacity) { |
123 | // Release all unused pages. |
124 | const size_t page_aligned_size = round_up_po2(size, xnn_params.page_size); |
125 | const uint8_t* mem_start = (uint8_t*) start; |
126 | const uint8_t* unused_start = mem_start + page_aligned_size; |
127 | assert(*capacity >= page_aligned_size); |
128 | const size_t unused_capacity = *capacity - page_aligned_size; |
129 | |
130 | xnn_log_debug("releasing memory, start %p, used: %zu, capacity: %zu, unused %zu" , mem_start, size, *capacity, |
131 | unused_capacity); |
132 | |
133 | if (unused_capacity != 0) { |
134 | // Free unused pages. |
135 | #if XNN_PLATFORM_WINDOWS |
136 | // We cannot selectively release pages inside the region of pages, so just decommit them. |
137 | if (!VirtualFree((void*) unused_start, unused_capacity, MEM_DECOMMIT)) { |
138 | xnn_log_error("failed to unmap code/weights buffer, error code: %" PRIu32, (uint32_t) GetLastError()); |
139 | return xnn_status_invalid_state; |
140 | } |
141 | *capacity = page_aligned_size; |
142 | #elif !XNN_PLATFORM_WEB |
143 | // Web does not support partial unmapping. |
144 | if (munmap((void*) unused_start, unused_capacity) == -1) { |
145 | xnn_log_error("failed to unmap code/weights buffer, error code: %d" , errno); |
146 | return xnn_status_invalid_state; |
147 | } |
148 | *capacity = page_aligned_size; |
149 | #else |
150 | if (unused_capacity == *capacity) { |
151 | if (munmap((void*) unused_start, unused_capacity) == -1) { |
152 | xnn_log_error("failed to unmap code/weights buffer, error code: %d" , errno); |
153 | return xnn_status_invalid_state; |
154 | } else { |
155 | *capacity = 0; |
156 | } |
157 | } |
158 | #endif |
159 | } |
160 | |
161 | return xnn_status_success; |
162 | } |
163 | |
164 | enum xnn_memory_permission { |
165 | xnn_memory_permission_read_only, |
166 | xnn_memory_permission_read_execute, |
167 | }; |
168 | |
169 | static enum xnn_status set_memory_permission(void* start, size_t size, enum xnn_memory_permission permission) { |
170 | #if XNN_PLATFORM_WINDOWS |
171 | DWORD old = 0, prot = 0; |
172 | switch (permission) { |
173 | case xnn_memory_permission_read_only: |
174 | prot = PAGE_READONLY; |
175 | break; |
176 | case xnn_memory_permission_read_execute: |
177 | prot = PAGE_EXECUTE_READ; |
178 | break; |
179 | default: |
180 | XNN_UNREACHABLE; |
181 | } |
182 | if (!VirtualProtect(start, size, prot, &old)) { |
183 | xnn_log_error( |
184 | "failed to set memory permission (%d), error code: %" PRIu32, permission, (uint32_t) GetLastError()); |
185 | return xnn_status_invalid_state; |
186 | } |
187 | #elif XNN_PLATFORM_WEB |
188 | // Memory protection not supported on Web. |
189 | return xnn_status_success; |
190 | #else |
191 | int prot = 0; |
192 | switch (permission) { |
193 | case xnn_memory_permission_read_only: |
194 | prot = PROT_READ; |
195 | break; |
196 | case xnn_memory_permission_read_execute: |
197 | prot = PROT_READ | PROT_EXEC; |
198 | break; |
199 | default: |
200 | XNN_UNREACHABLE; |
201 | } |
202 | if (mprotect(start, size, prot) == -1) { |
203 | xnn_log_error("failed to set memory permission (%d), error code: %d" , permission, errno); |
204 | return xnn_status_invalid_state; |
205 | } |
206 | #endif |
207 | return xnn_status_success; |
208 | } |
209 | |
210 | #if XNN_PLATFORM_JIT |
211 | enum xnn_status xnn_finalize_code_memory(struct xnn_code_buffer* buf) { |
212 | enum xnn_status status; |
213 | status = release_unused_memory(buf->size, buf->start, &buf->capacity); |
214 | if (status != xnn_status_success) { |
215 | return status; |
216 | } |
217 | |
218 | if (buf->capacity == 0) { |
219 | return xnn_status_success; |
220 | } |
221 | |
222 | // Flush icache, do it before changing permissions due to bugs on older ARM64 kernels. |
223 | #if (XNN_ARCH_ARM || XNN_ARCH_ARM64) && XNN_PLATFORM_JIT |
224 | // iOS toolchain doesn't support this, use sys_icache_invalidate, when we support iOS. |
225 | __builtin___clear_cache(buf->start, (void*) ((uint8_t*) buf->start + buf->capacity)); |
226 | #endif // (XNN_ARCH_ARM || XNN_ARCH_ARM64) && !XNN_PLATFORM_IOS |
227 | |
228 | // Set permissions to RX (no write). |
229 | #if XNN_PLATFORM_WINDOWS |
230 | DWORD old = 0; |
231 | if (!VirtualProtect(buf->start, buf->size, PAGE_EXECUTE_READ, &old)) { |
232 | xnn_log_error("failed to make code buffer read+execute, error code: %" PRIu32, (uint32_t) GetLastError()); |
233 | return xnn_status_invalid_state; |
234 | } |
235 | #else |
236 | if (mprotect(buf->start, buf->size, PROT_READ | PROT_EXEC) == -1) { |
237 | xnn_log_error("failed to make code buffer read+execute, error code: %d" , errno); |
238 | return xnn_status_invalid_state; |
239 | } |
240 | #endif |
241 | return set_memory_permission(buf->start, buf->size, xnn_memory_permission_read_execute); |
242 | } |
243 | #endif // XNN_PLATFORM_JIT |
244 | |
245 | enum xnn_status xnn_release_code_memory(struct xnn_code_buffer* buf) { |
246 | if (buf->capacity == 0) { |
247 | return xnn_status_success; |
248 | } |
249 | const enum xnn_status status = release_memory(buf->start, buf->capacity); |
250 | if (status != xnn_status_success) { |
251 | return status; |
252 | } |
253 | memset(buf, 0, sizeof(struct xnn_code_buffer)); |
254 | return xnn_status_success; |
255 | } |
256 | |
257 | enum xnn_status xnn_reserve_code_memory(struct xnn_code_buffer* buf, size_t n) { |
258 | if (buf->size + n <= buf->capacity) { |
259 | return xnn_status_success; |
260 | } |
261 | xnn_log_debug("reserving code memory of size %zu" , n); |
262 | |
263 | size_t new_capacity = 0; |
264 | void* new_start = resize_buffer(buf->start, buf->size, buf->capacity, buf->size + n, &new_capacity); |
265 | if (new_start == NULL) { |
266 | xnn_log_error("failed to reserve code memory" ); |
267 | return xnn_status_out_of_memory; |
268 | } |
269 | buf->start = new_start; |
270 | buf->capacity = new_capacity; |
271 | return xnn_status_success; |
272 | } |
273 | |
274 | enum xnn_status xnn_allocate_weights_memory(struct xnn_weights_buffer* buf, size_t size) { |
275 | memset(buf, 0, sizeof(struct xnn_weights_buffer)); |
276 | size_t page_aligned_size = round_up_po2(size, xnn_params.page_size); |
277 | buf->start = allocate_buffer(page_aligned_size); |
278 | if (buf->start == NULL) { |
279 | return xnn_status_out_of_memory; |
280 | } |
281 | |
282 | buf->size = 0; |
283 | buf->capacity = page_aligned_size; |
284 | return xnn_status_success; |
285 | } |
286 | |
287 | enum xnn_status xnn_release_weights_memory(struct xnn_weights_buffer* buf) { |
288 | if (buf->capacity == 0) { |
289 | return xnn_status_success; |
290 | } |
291 | enum xnn_status status = release_memory(buf->start, buf->capacity); |
292 | if (status != xnn_status_success) { |
293 | return status; |
294 | } |
295 | memset(buf, 0, sizeof(struct xnn_code_buffer)); |
296 | return xnn_status_success; |
297 | } |
298 | |
299 | enum xnn_status xnn_reserve_weights_memory(struct xnn_weights_buffer* buf, size_t n) { |
300 | if (buf->size + n <= buf->capacity) { |
301 | xnn_log_debug("reserving weights memory of size %zu without growing buffer" , n); |
302 | return xnn_status_success; |
303 | } |
304 | |
305 | size_t new_capacity = 0; |
306 | void* new_start = resize_buffer(buf->start, buf->size, buf->capacity, buf->size + n, &new_capacity); |
307 | if (new_start == NULL) { |
308 | xnn_log_error("failed to reserve weights memory" ); |
309 | return xnn_status_out_of_memory; |
310 | } |
311 | buf->start = new_start; |
312 | buf->capacity = new_capacity; |
313 | |
314 | return xnn_status_success; |
315 | } |
316 | |
317 | enum xnn_status xnn_finalize_weights_memory(struct xnn_weights_buffer* buf) { |
318 | enum xnn_status status; |
319 | status = release_unused_memory(buf->size, buf->start, &buf->capacity); |
320 | if (status != xnn_status_success) { |
321 | return status; |
322 | } |
323 | |
324 | if (buf->capacity == 0) { |
325 | return xnn_status_success; |
326 | } |
327 | |
328 | return set_memory_permission(buf->start, buf->size, xnn_memory_permission_read_only); |
329 | } |
330 | |