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 | |
35 | static size_t system_page_size = 0; |
36 | |
37 | static 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. |
59 | static 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. |
84 | static 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. |
104 | static 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 | |
134 | enum 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`. |
148 | static 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 | |
190 | enum xnn_memory_permission { |
191 | xnn_memory_permission_read_only, |
192 | xnn_memory_permission_read_execute, |
193 | }; |
194 | |
195 | static 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 |
237 | enum 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 | |
274 | enum 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 | |
286 | enum 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 | |
304 | enum 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 | |
317 | enum 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 | |
329 | enum 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 | |
348 | enum 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 | |