blob: 252f9758888e60f684004322d91ec1a902a5e0d4 [file] [log] [blame] [edit]
// Copyright 2021 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.
// Include first for the platform detection macros.
#include "xnnpack/common.h"
#if XNN_PLATFORM_WINDOWS
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <windows.h>
#else
// This define needs to come first because errno include features.h and would have defined macros that lead to
// sys/mman.h not having mremap.
#if !defined(_GNU_SOURCE)
#define _GNU_SOURCE
#endif
#include <errno.h>
#include <sys/mman.h>
#include <unistd.h>
#endif
#include <stddef.h>
#include <stdint.h>
#include <xnnpack.h>
#include "xnnpack/allocator.h"
#include "xnnpack/log.h"
#include "xnnpack/math.h"
#include "xnnpack/params.h"
// Helpers to allocate/mmap and release memory used by both code and weights cache.
// Maps `size` bytes of memory, returns pointer to allocation, NULL if failed.
static void* allocate_buffer(size_t size) {
xnn_log_debug("allocating buffer of size %zu", size);
assert(size == round_up_po2(size, xnn_params.page_size));
#if XNN_PLATFORM_WINDOWS
void* p = VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
if (p == NULL) {
xnn_log_error("failed to allocate %zu bytes for code/weights buffer, error code: %" PRIu32,
size, (uint32_t) GetLastError());
return NULL;
}
#else
void* p = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
if (p == MAP_FAILED) {
xnn_log_error("failed to allocate %zu bytes for code/weights buffer, error code: %d", size, errno);
return NULL;
}
#endif
return p;
}
// Releases memory previously mapped by `allocate_buffer`, returns xnn_status_success on success.
static enum xnn_status release_memory(void* start, size_t capacity) {
#if XNN_PLATFORM_WINDOWS
// We only decommited any unused capacity, so we release all of it now.
if (!VirtualFree(start, 0, MEM_RELEASE)) {
xnn_log_error("failed to release code/weights buffer, error code: %" PRIu32, (uint32_t) GetLastError());
return xnn_status_invalid_state;
}
#else
if (munmap(start, capacity) == -1) {
xnn_log_error("failed to release code/weights buffer, error code: %d", errno);
return xnn_status_invalid_state;
}
#endif
return xnn_status_success;
}
// Resize a buffer at old_pointer of size old_bytes to new_size. The actual new size of the resized buffer is written to
// new_capacity_out, which can be >= new_size due to page alignment requirements.
// Returns a pointer to a buffer which might be the same as old_pointer if we can remap virtual memory, otherwise we
// allocate a new buffer and copy contents of old_buffer over.
static void* resize_buffer(
void* old_pointer, size_t old_size, size_t old_capacity, size_t new_size, size_t* new_capacity_out)
{
size_t new_capacity = round_up_po2(new_size, xnn_params.page_size);
#if XNN_PLATFORM_LINUX
void* new_pointer = mremap(old_pointer, old_size, new_capacity, MREMAP_MAYMOVE, NULL);
if (new_pointer == MAP_FAILED) {
xnn_log_error("mremap failed with errno: %d", errno);
return NULL;
}
xnn_log_debug("resize_buffer: remap, old capacity %zu to new capacity %zu", old_capacity, new_capacity);
#else
void* new_pointer = allocate_buffer(new_capacity);
if (new_pointer == NULL) {
xnn_log_error("allocate_buffer failed");
return NULL;
}
memcpy(new_pointer, old_pointer, old_size);
// Release old code_buffer.
enum xnn_status status = release_memory(old_pointer, old_capacity);
if (status != xnn_status_success) {
xnn_log_error("releasing old buffer failed, this could be a leak of %zu bytes", old_capacity);
// Log but proceed as per normal since we successfully allocated a new memory that can be used by the caller.
}
xnn_log_debug("resize_buffer: allocate memory, old capacity %zu to new capacity %zu", old_capacity, new_capacity);
#endif
*new_capacity_out = new_capacity;
return new_pointer;
}
enum xnn_status xnn_allocate_code_memory(struct xnn_code_buffer* buf, size_t size) {
memset(buf, 0, sizeof(struct xnn_code_buffer));
size_t page_aligned_size = round_up_po2(size, xnn_params.page_size);
buf->start = allocate_buffer(page_aligned_size);
if (buf->start == NULL) {
return xnn_status_out_of_memory;
}
buf->size = 0;
buf->capacity = page_aligned_size;
return xnn_status_success;
}
// Releases unused memory. Will write the new capacity to `capacity`.
static enum xnn_status release_unused_memory(size_t size, void* start, size_t* capacity) {
// Release all unused pages.
const size_t page_aligned_size = round_up_po2(size, xnn_params.page_size);
const uint8_t* mem_start = (uint8_t*) start;
const uint8_t* unused_start = mem_start + page_aligned_size;
assert(*capacity >= page_aligned_size);
const size_t unused_capacity = *capacity - page_aligned_size;
xnn_log_debug("releasing memory, start %p, used: %zu, capacity: %zu, unused %zu", mem_start, size, *capacity,
unused_capacity);
if (unused_capacity != 0) {
// Free unused pages.
#if XNN_PLATFORM_WINDOWS
// We cannot selectively release pages inside the region of pages, so just decommit them.
if (!VirtualFree((void*) unused_start, unused_capacity, MEM_DECOMMIT)) {
xnn_log_error("failed to unmap code/weights buffer, error code: %" PRIu32, (uint32_t) GetLastError());
return xnn_status_invalid_state;
}
*capacity = page_aligned_size;
#elif !XNN_PLATFORM_WEB
// Web does not support partial unmapping.
if (munmap((void*) unused_start, unused_capacity) == -1) {
xnn_log_error("failed to unmap code/weights buffer, error code: %d", errno);
return xnn_status_invalid_state;
}
*capacity = page_aligned_size;
#else
if (unused_capacity == *capacity) {
if (munmap((void*) unused_start, unused_capacity) == -1) {
xnn_log_error("failed to unmap code/weights buffer, error code: %d", errno);
return xnn_status_invalid_state;
} else {
*capacity = 0;
}
}
#endif
}
return xnn_status_success;
}
enum xnn_memory_permission {
xnn_memory_permission_read_only,
xnn_memory_permission_read_execute,
};
static enum xnn_status set_memory_permission(void* start, size_t size, enum xnn_memory_permission permission) {
#if XNN_PLATFORM_WINDOWS
DWORD old = 0, prot = 0;
switch (permission) {
case xnn_memory_permission_read_only:
prot = PAGE_READONLY;
break;
case xnn_memory_permission_read_execute:
prot = PAGE_EXECUTE_READ;
break;
default:
XNN_UNREACHABLE;
}
if (!VirtualProtect(start, size, prot, &old)) {
xnn_log_error(
"failed to set memory permission (%d), error code: %" PRIu32, permission, (uint32_t) GetLastError());
return xnn_status_invalid_state;
}
#elif XNN_PLATFORM_WEB
// Memory protection not supported on Web.
return xnn_status_success;
#else
int prot = 0;
switch (permission) {
case xnn_memory_permission_read_only:
prot = PROT_READ;
break;
case xnn_memory_permission_read_execute:
prot = PROT_READ | PROT_EXEC;
break;
default:
XNN_UNREACHABLE;
}
if (mprotect(start, size, prot) == -1) {
xnn_log_error("failed to set memory permission (%d), error code: %d", permission, errno);
return xnn_status_invalid_state;
}
#endif
return xnn_status_success;
}
#if XNN_PLATFORM_JIT
enum xnn_status xnn_finalize_code_memory(struct xnn_code_buffer* buf) {
enum xnn_status status;
status = release_unused_memory(buf->size, buf->start, &buf->capacity);
if (status != xnn_status_success) {
return status;
}
if (buf->capacity == 0) {
return xnn_status_success;
}
// Flush icache, do it before changing permissions due to bugs on older ARM64 kernels.
#if (XNN_ARCH_ARM || XNN_ARCH_ARM64) && XNN_PLATFORM_JIT
// iOS toolchain doesn't support this, use sys_icache_invalidate, when we support iOS.
__builtin___clear_cache(buf->start, (void*) ((uint8_t*) buf->start + buf->capacity));
#endif // (XNN_ARCH_ARM || XNN_ARCH_ARM64) && !XNN_PLATFORM_IOS
// Set permissions to RX (no write).
#if XNN_PLATFORM_WINDOWS
DWORD old = 0;
if (!VirtualProtect(buf->start, buf->size, PAGE_EXECUTE_READ, &old)) {
xnn_log_error("failed to make code buffer read+execute, error code: %" PRIu32, (uint32_t) GetLastError());
return xnn_status_invalid_state;
}
#else
if (mprotect(buf->start, buf->size, PROT_READ | PROT_EXEC) == -1) {
xnn_log_error("failed to make code buffer read+execute, error code: %d", errno);
return xnn_status_invalid_state;
}
#endif
return set_memory_permission(buf->start, buf->size, xnn_memory_permission_read_execute);
}
#endif // XNN_PLATFORM_JIT
enum xnn_status xnn_release_code_memory(struct xnn_code_buffer* buf) {
if (buf->capacity == 0) {
return xnn_status_success;
}
const enum xnn_status status = release_memory(buf->start, buf->capacity);
if (status != xnn_status_success) {
return status;
}
memset(buf, 0, sizeof(struct xnn_code_buffer));
return xnn_status_success;
}
enum xnn_status xnn_reserve_code_memory(struct xnn_code_buffer* buf, size_t n) {
if (buf->size + n <= buf->capacity) {
return xnn_status_success;
}
xnn_log_debug("reserving code memory of size %zu", n);
size_t new_capacity = 0;
void* new_start = resize_buffer(buf->start, buf->size, buf->capacity, buf->size + n, &new_capacity);
if (new_start == NULL) {
xnn_log_error("failed to reserve code memory");
return xnn_status_out_of_memory;
}
buf->start = new_start;
buf->capacity = new_capacity;
return xnn_status_success;
}
enum xnn_status xnn_allocate_weights_memory(struct xnn_weights_buffer* buf, size_t size) {
memset(buf, 0, sizeof(struct xnn_weights_buffer));
size_t page_aligned_size = round_up_po2(size, xnn_params.page_size);
buf->start = allocate_buffer(page_aligned_size);
if (buf->start == NULL) {
return xnn_status_out_of_memory;
}
buf->size = 0;
buf->capacity = page_aligned_size;
return xnn_status_success;
}
enum xnn_status xnn_release_weights_memory(struct xnn_weights_buffer* buf) {
if (buf->capacity == 0) {
return xnn_status_success;
}
enum xnn_status status = release_memory(buf->start, buf->capacity);
if (status != xnn_status_success) {
return status;
}
memset(buf, 0, sizeof(struct xnn_code_buffer));
return xnn_status_success;
}
enum xnn_status xnn_reserve_weights_memory(struct xnn_weights_buffer* buf, size_t n) {
if (buf->size + n <= buf->capacity) {
xnn_log_debug("reserving weights memory of size %zu without growing buffer", n);
return xnn_status_success;
}
size_t new_capacity = 0;
void* new_start = resize_buffer(buf->start, buf->size, buf->capacity, buf->size + n, &new_capacity);
if (new_start == NULL) {
xnn_log_error("failed to reserve weights memory");
return xnn_status_out_of_memory;
}
buf->start = new_start;
buf->capacity = new_capacity;
return xnn_status_success;
}
enum xnn_status xnn_finalize_weights_memory(struct xnn_weights_buffer* buf) {
enum xnn_status status;
status = release_unused_memory(buf->size, buf->start, &buf->capacity);
if (status != xnn_status_success) {
return status;
}
if (buf->capacity == 0) {
return xnn_status_success;
}
return set_memory_permission(buf->start, buf->size, xnn_memory_permission_read_only);
}