Support overriding memory allocation functions
- Let users provide their own memory management functions for XNNPACK
PiperOrigin-RevId: 281355722
diff --git a/src/xnnpack/allocator.h b/src/xnnpack/allocator.h
index c946656..8aef8b2 100644
--- a/src/xnnpack/allocator.h
+++ b/src/xnnpack/allocator.h
@@ -8,40 +8,42 @@
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
-#ifdef __ANDROID__
- #include <malloc.h>
-#endif
#include <xnnpack/common.h>
-
-extern int posix_memalign(void **memptr, size_t alignment, size_t size);
+#include <xnnpack/params.h>
-#define XNN_ALLOCATION_ALIGNMENT 16
-
+#if XNN_ARCH_ASMJS || XNN_ARCH_WASM
+ #define XNN_ALLOCATION_ALIGNMENT 4
+#else
+ #define XNN_ALLOCATION_ALIGNMENT 16
+#endif
inline static void* xnn_allocate_memory(size_t memory_size) {
- void* memory_ptr = NULL;
-#if XNN_ARCH_ASMJS || XNN_ARCH_WASM
- memory_ptr = malloc(memory_size);
-#elif defined(__ANDROID__)
- memory_ptr = memalign(XNN_ALLOCATION_ALIGNMENT, memory_size);
-#else
- if (posix_memalign(&memory_ptr, XNN_ALLOCATION_ALIGNMENT, memory_size) != 0) {
- return NULL;
- }
-#endif
- return memory_ptr;
+ return xnn_params.allocator.allocate(xnn_params.allocator.context, memory_size);
}
-inline static void* xnn_allocate_zero_memory(size_t memory_size) {
- void* memory_ptr = xnn_allocate_memory(memory_size);
- if (memory_ptr != NULL) {
- memset(memory_ptr, 0, memory_size);
- }
- return memory_ptr;
+inline static void* xnn_reallocate_memory(void* memory_pointer, size_t memory_size) {
+ return xnn_params.allocator.reallocate(xnn_params.allocator.context, memory_pointer, memory_size);
}
-inline static void xnn_release_memory(void* memory_ptr) {
- free(memory_ptr);
+inline static void xnn_release_memory(void* memory_pointer) {
+ xnn_params.allocator.deallocate(xnn_params.allocator.context, memory_pointer);
+}
+
+inline static void* xnn_allocate_simd_memory(size_t memory_size) {
+ return xnn_params.allocator.aligned_allocate(xnn_params.allocator.context, XNN_ALLOCATION_ALIGNMENT, memory_size);
+}
+
+inline static void* xnn_allocate_zero_simd_memory(size_t memory_size) {
+ void* memory_pointer = xnn_params.allocator.aligned_allocate(
+ xnn_params.allocator.context, XNN_ALLOCATION_ALIGNMENT, memory_size);
+ if (memory_pointer != NULL) {
+ memset(memory_pointer, 0, memory_size);
+ }
+ return memory_pointer;
+}
+
+inline static void xnn_release_simd_memory(void* memory_pointer) {
+ xnn_params.allocator.aligned_deallocate(xnn_params.allocator.context, memory_pointer);
}