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