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 #define WIN32_LEAN_AND_MEAN
11 #include <windows.h>
12 #else
13 #include <errno.h>
14 #include <sys/mman.h>
15 #include <unistd.h>
16 #endif
17
18 #include <stddef.h>
19 #include <stdint.h>
20 #include <xnnpack.h>
21
22 #include "xnnpack/allocator.h"
23 #include "xnnpack/log.h"
24 #include "xnnpack/math.h"
25
xnn_allocate_code_memory(struct xnn_code_buffer * buf,size_t size)26 enum xnn_status xnn_allocate_code_memory(struct xnn_code_buffer* buf, size_t size) {
27 #if XNN_PLATFORM_WINDOWS
28 void* p = VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
29 if (p == NULL) {
30 xnn_log_error("failed to allocate %zu bytes for JIT code buffer, error code: %" PRIu32,
31 size, (uint32_t) GetLastError());
32 return xnn_status_out_of_memory;
33 }
34 #else
35 void* p = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
36 if (p == MAP_FAILED) {
37 xnn_log_error("failed to allocate %zu bytes for JIT code buffer, error code: %d", size, errno);
38 return xnn_status_out_of_memory;
39 }
40 #endif
41
42 buf->code = p;
43 buf->size = 0;
44 buf->capacity = size;
45 return xnn_status_success;
46 }
47
xnn_finalize_code_memory(struct xnn_code_buffer * buf)48 enum xnn_status xnn_finalize_code_memory(struct xnn_code_buffer* buf) {
49 // Get page size.
50 #if XNN_PLATFORM_WINDOWS
51 SYSTEM_INFO sysinfo;
52 GetSystemInfo(&sysinfo);
53 const size_t page_size = sysinfo.dwPageSize;
54 #else
55 const long res = sysconf(_SC_PAGESIZE);
56 if (res == -1) {
57 xnn_log_error("failed to get page size, error code: %d", errno);
58 return xnn_status_invalid_state;
59 }
60 const size_t page_size = res;
61 #endif
62
63 // Release all unused pages.
64 const size_t page_aligned_code_size = round_up_po2(buf->size, page_size);
65 const uint8_t* start = (uint8_t*) buf->code;
66 const uint8_t* unused_start = start + page_aligned_code_size;
67 const size_t unused_capacity = buf->capacity - page_aligned_code_size;
68
69 xnn_log_debug("JIT code memory start %p, used: %zu, capacity: %zu, unused %zu", start, buf->size, buf->capacity,
70 unused_capacity);
71
72 if (unused_capacity != 0) {
73 // Free unused pages.
74 #if XNN_PLATFORM_WINDOWS
75 // We cannot selectively release pages inside the region of pages, so just decommit them.
76 if (!VirtualFree((void*) unused_start, unused_capacity, MEM_DECOMMIT)) {
77 xnn_log_error("failed to unmap code buffer, error code: %" PRIu32, (uint32_t) GetLastError());
78 return xnn_status_invalid_state;
79 }
80 #else
81 if (munmap((void*) unused_start, unused_capacity) == -1) {
82 xnn_log_error("failed to unmap code buffer, error code: %d", errno);
83 return xnn_status_invalid_state;
84 }
85 #endif
86 }
87
88 buf->capacity = page_aligned_code_size;
89
90 if (buf->capacity == 0) {
91 return xnn_status_success;
92 }
93
94 // Flush icache, do it before changing permissions due to bugs on older ARM64 kernels.
95 #if (XNN_ARCH_ARM || XNN_ARCH_ARM64) && !XNN_PLATFORM_IOS
96 // iOS toolchain doesn't support this, use sys_icache_invalidate, when we support iOS.
97 __builtin___clear_cache(buf->code, (void*) ((uint8_t*) buf->code + buf->capacity));
98 #endif // (XNN_ARCH_ARM || XNN_ARCH_ARM64) && !XNN_PLATFORM_IOS
99
100 // Set permissions to RX (no write).
101 #if XNN_PLATFORM_WINDOWS
102 DWORD old = 0;
103 if (!VirtualProtect(buf->code, buf->size, PAGE_EXECUTE_READ, &old)) {
104 xnn_log_error("failed to make code buffer read+execute, error code: %" PRIu32, (uint32_t) GetLastError());
105 return xnn_status_invalid_state;
106 }
107 #else
108 if (mprotect(buf->code, buf->size, PROT_READ | PROT_EXEC) == -1) {
109 xnn_log_error("failed to make code buffer read+execute, error code: %d", errno);
110 return xnn_status_invalid_state;
111 }
112 #endif
113 return xnn_status_success;
114 }
115
xnn_release_code_memory(struct xnn_code_buffer * buf)116 enum xnn_status xnn_release_code_memory(struct xnn_code_buffer* buf) {
117 if (buf->capacity == 0) {
118 return xnn_status_success;
119 }
120 #if XNN_PLATFORM_WINDOWS
121 // We only decommited any unused capacity, so we release all of it now.
122 if (!VirtualFree(buf->code, 0, MEM_RELEASE)) {
123 xnn_log_error("failed to release code buffer for JIT, error code: %" PRIu32, (uint32_t) GetLastError());
124 return xnn_status_invalid_state;
125 }
126 #else
127 if (munmap(buf->code, buf->capacity) == -1) {
128 xnn_log_error("failed to release code buffer for JIT, error code: %d", errno);
129 return xnn_status_invalid_state;
130 }
131 #endif
132 return xnn_status_success;
133 }
134