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.
37static 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.
58static 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.
78static 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
108enum 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`.
122static 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
164enum xnn_memory_permission {
165 xnn_memory_permission_read_only,
166 xnn_memory_permission_read_execute,
167};
168
169static 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
211enum 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
245enum 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
257enum 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
274enum 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
287enum 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
299enum 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
317enum 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