Add abstraction on top of ba_alloc which supports arbitrary size allocations

src/base_alloc/base_alloc_global.c

@@ -5,49 +5,116 @@
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 */
 
+/* A MT-safe base allocator */
+
 #include <assert.h>
+#include <stdio.h>
 #include <stdlib.h>
 
 #include "base_alloc.h"
 #include "base_alloc_global.h"
+#include "base_alloc_internal.h"
 #include "utils_concurrency.h"
-
-#define SIZE_BA_POOL_CHUNK 128
+#include "utils_math.h"
 
 // global base allocator used by all providers and pools
-static umf_ba_pool_t *BA_pool = NULL;
 static UTIL_ONCE_FLAG ba_is_initialized = UTIL_ONCE_FLAG_INIT;
 
+// allocation classes need to be powers of 2
+#define ALLOCATION_CLASSES                                                     \
+    { 16, 32, 64, 128 }
+#define NUM_ALLOCATION_CLASSES 4
+
+struct base_alloc_t {
+    size_t ac_sizes[NUM_ALLOCATION_CLASSES];
+    umf_ba_pool_t *ac[NUM_ALLOCATION_CLASSES];
+    size_t smallest_ac_size_log2;
+};
+
+static struct base_alloc_t BASE_ALLOC = {.ac_sizes = ALLOCATION_CLASSES};
+
+void umf_ba_destroy_global(void) {
+    for (int i = 0; i < NUM_ALLOCATION_CLASSES; i++) {
+        if (BASE_ALLOC.ac[i]) {
+            umf_ba_destroy(BASE_ALLOC.ac[i]);
+            BASE_ALLOC.ac[i] = NULL;
+        }
+    }
+}
+
 static void umf_ba_create_global(void) {
-    assert(BA_pool == NULL);
-    BA_pool = umf_ba_create(SIZE_BA_POOL_CHUNK);
-    assert(BA_pool);
+    for (int i = 0; i < NUM_ALLOCATION_CLASSES; i++) {
+        // allocation classes need to be powers of 2
+        assert(0 == (BASE_ALLOC.ac_sizes[i] & (BASE_ALLOC.ac_sizes[i] - 1)));
+        BASE_ALLOC.ac[i] = umf_ba_create(BASE_ALLOC.ac_sizes[i]);
+        if (!BASE_ALLOC.ac[i]) {
+            fprintf(stderr,
+                    "base_alloc: Error. Cannot create base alloc allocation "
+                    "class for size: %zu\n. Each allocation will fallback to "
+                    "allocating memory from the OS.",
+                    BASE_ALLOC.ac_sizes[i]);
+        }
+    }
+
+    size_t smallestSize = BASE_ALLOC.ac_sizes[0];
+    BASE_ALLOC.smallest_ac_size_log2 = log2Utils(smallestSize);
+
 #if defined(_WIN32) && !defined(UMF_SHARED_LIBRARY)
     atexit(umf_ba_destroy_global);
 #endif
 }
 
-void umf_ba_destroy_global(void) {
-    if (BA_pool) {
-        umf_ba_pool_t *pool = BA_pool;
-        BA_pool = NULL;
-        umf_ba_destroy(pool);
+// returns index of the allocation class for a given size
+static int size_to_idx(size_t size) {
+    assert(size <= BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]);
+
+    if (size <= BASE_ALLOC.ac_sizes[0]) {
+        return 0;
     }
+
+    int isPowerOf2 = (0 == (size & (size - 1)));
+    int index =
+        (int)(log2Utils(size) + !isPowerOf2 - BASE_ALLOC.smallest_ac_size_log2);
+
+    assert(index >= 0);
+    assert(index < NUM_ALLOCATION_CLASSES);
+
+    return index;
 }
 
-umf_ba_pool_t *umf_ba_get_pool(size_t size) {
+void *umf_ba_global_alloc(size_t size) {
     util_init_once(&ba_is_initialized, umf_ba_create_global);
 
-    if (!BA_pool) {
-        return NULL;
+    if (size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
+        fprintf(stderr,
+                "base_alloc: allocation size larger than the biggest "
+                "allocation class. Falling back to OS memory allocation.\n");
+        return ba_os_alloc(size);
     }
 
-    // TODO: a specific class-size base allocator can be returned here
-    assert(size <= SIZE_BA_POOL_CHUNK);
+    int ac_index = size_to_idx(size);
+    if (!BASE_ALLOC.ac[ac_index]) {
+        // if creating ac failed, fall back to os allocation
+        fprintf(stderr, "base_alloc: allocation class not created. Falling "
+                        "back to OS memory allocation.\n");
+        return ba_os_alloc(size);
+    }
+
+    return umf_ba_alloc(BASE_ALLOC.ac[ac_index]);
+}
+
+void umf_ba_global_free(void *ptr, size_t size) {
+    if (size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
+        ba_os_free(ptr, size);
+        return;
+    }
 
-    if (size > SIZE_BA_POOL_CHUNK) {
-        return NULL;
+    int ac_index = size_to_idx(size);
+    if (!BASE_ALLOC.ac[ac_index]) {
+        // if creating ac failed, memory must have been allocated by os
+        ba_os_free(ptr, size);
+        return;
     }
 
-    return BA_pool;
+    umf_ba_free(BASE_ALLOC.ac[ac_index], ptr);
 }

src/base_alloc/base_alloc_global.h

@@ -14,7 +14,8 @@
 extern "C" {
 #endif
 
-umf_ba_pool_t *umf_ba_get_pool(size_t size);
+void *umf_ba_global_alloc(size_t size);
+void umf_ba_global_free(void *ptr, size_t size);
 void umf_ba_destroy_global(void);
 
 #ifdef __cplusplus

src/memory_pool_default.c

@@ -25,26 +25,21 @@ umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
     }
 
     umf_result_t ret = UMF_RESULT_SUCCESS;
-    umf_ba_pool_t *base_allocator = umf_ba_get_pool(sizeof(umf_memory_pool_t));
-    if (!base_allocator) {
-        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
-    }
-
-    umf_memory_pool_handle_t pool = umf_ba_alloc(base_allocator);
+    umf_memory_pool_handle_t pool =
+        umf_ba_global_alloc(sizeof(umf_memory_pool_t));
     if (!pool) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
     assert(ops->version == UMF_VERSION_CURRENT);
 
-    pool->base_allocator = base_allocator;
     pool->provider = provider;
     pool->own_provider = false;
 
     pool->ops = *ops;
     ret = ops->initialize(pool->provider, params, &pool->pool_priv);
     if (ret != UMF_RESULT_SUCCESS) {
-        umf_ba_free(base_allocator, pool);
+        umf_ba_global_free(pool, sizeof(umf_memory_pool_t));
         return ret;
     }
 
@@ -61,7 +56,7 @@ void umfPoolDestroy(umf_memory_pool_handle_t hPool) {
         umfMemoryProviderDestroy(hProvider);
     }
     // TODO: this free keeps memory in base allocator, so it can lead to OOM in some scenarios (it should be optimized)
-    umf_ba_free(hPool->base_allocator, hPool);
+    umf_ba_global_free(hPool, sizeof(umf_memory_pool_t));
 }
 
 umf_result_t umfFree(void *ptr) {

src/memory_pool_internal.h

@@ -31,9 +31,6 @@ typedef struct umf_memory_pool_t {
     umf_memory_provider_handle_t provider;
     // Tells whether memory provider is owned by the pool.
     bool own_provider;
-
-    // saved pointer to the global base allocator
-    umf_ba_pool_t *base_allocator;
 } umf_memory_pool_t;
 
 umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,

src/memory_pool_tracking.c

@@ -26,12 +26,8 @@ umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
     }
 
     umf_result_t ret = UMF_RESULT_SUCCESS;
-    umf_ba_pool_t *base_allocator = umf_ba_get_pool(sizeof(umf_memory_pool_t));
-    if (!base_allocator) {
-        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
-    }
-
-    umf_memory_pool_handle_t pool = umf_ba_alloc(base_allocator);
+    umf_memory_pool_handle_t pool =
+        umf_ba_global_alloc(sizeof(umf_memory_pool_t));
     if (!pool) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
@@ -44,7 +40,6 @@ umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
         goto err_provider_create;
     }
 
-    pool->base_allocator = base_allocator;
     pool->own_provider = false;
 
     pool->ops = *ops;
@@ -59,7 +54,7 @@ umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
 err_pool_init:
     umfMemoryProviderDestroy(pool->provider);
 err_provider_create:
-    umf_ba_free(base_allocator, pool);
+    umf_ba_global_free(pool, sizeof(umf_memory_pool_t));
     return ret;
 }
 
@@ -74,7 +69,7 @@ void umfPoolDestroy(umf_memory_pool_handle_t hPool) {
     // Destroy tracking provider.
     umfMemoryProviderDestroy(hPool->provider);
     // TODO: this free keeps memory in base allocator, so it can lead to OOM in some scenarios (it should be optimized)
-    umf_ba_free(hPool->base_allocator, hPool);
+    umf_ba_global_free(hPool, sizeof(umf_memory_pool_t));
 }
 
 umf_result_t umfFree(void *ptr) {

src/memory_provider.c

@@ -22,9 +22,6 @@
 typedef struct umf_memory_provider_t {
     umf_memory_provider_ops_t ops;
     void *provider_priv;
-
-    // saved pointer to the global base allocator
-    umf_ba_pool_t *base_allocator;
 } umf_memory_provider_t;
 
 umf_result_t umfMemoryProviderCreate(const umf_memory_provider_ops_t *ops,
@@ -35,26 +32,20 @@ umf_result_t umfMemoryProviderCreate(const umf_memory_provider_ops_t *ops,
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    umf_ba_pool_t *base_allocator =
-        umf_ba_get_pool(sizeof(umf_memory_provider_t));
-    if (!base_allocator) {
-        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
-    }
-
-    umf_memory_provider_handle_t provider = umf_ba_alloc(base_allocator);
+    umf_memory_provider_handle_t provider =
+        umf_ba_global_alloc(sizeof(umf_memory_provider_t));
     if (!provider) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
     assert(ops->version == UMF_VERSION_CURRENT);
 
     provider->ops = *ops;
-    provider->base_allocator = base_allocator;
 
     void *provider_priv;
     umf_result_t ret = ops->initialize(params, &provider_priv);
     if (ret != UMF_RESULT_SUCCESS) {
-        umf_ba_free(base_allocator, provider);
+        umf_ba_global_free(provider, sizeof(umf_memory_provider_t));
         return ret;
     }
 
@@ -67,7 +58,7 @@ umf_result_t umfMemoryProviderCreate(const umf_memory_provider_ops_t *ops,
 
 void umfMemoryProviderDestroy(umf_memory_provider_handle_t hProvider) {
     hProvider->ops.finalize(hProvider->provider_priv);
-    umf_ba_free(hProvider->base_allocator, hProvider);
+    umf_ba_global_free(hProvider, sizeof(umf_memory_provider_t));
 }
 
 static void

src/memory_target.c

@@ -24,27 +24,20 @@ umf_result_t umfMemoryTargetCreate(const umf_memory_target_ops_t *ops,
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    umf_ba_pool_t *base_allocator =
-        umf_ba_get_pool(sizeof(umf_memory_target_t));
-    if (!base_allocator) {
-        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
-    }
-
     umf_memory_target_handle_t target =
-        (umf_memory_target_t *)umf_ba_alloc(base_allocator);
+        (umf_memory_target_t *)umf_ba_global_alloc(sizeof(umf_memory_target_t));
     if (!target) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
     assert(ops->version == UMF_VERSION_CURRENT);
 
-    target->base_allocator = base_allocator;
     target->ops = ops;
 
     void *target_priv;
     umf_result_t ret = ops->initialize(params, &target_priv);
     if (ret != UMF_RESULT_SUCCESS) {
-        umf_ba_free(base_allocator, target);
+        umf_ba_global_free(target, sizeof(umf_memory_target_t));
         return ret;
     }
 
@@ -58,5 +51,5 @@ umf_result_t umfMemoryTargetCreate(const umf_memory_target_ops_t *ops,
 void umfMemoryTargetDestroy(umf_memory_target_handle_t memoryTarget) {
     assert(memoryTarget);
     memoryTarget->ops->finalize(memoryTarget->priv);
-    umf_ba_free(memoryTarget->base_allocator, memoryTarget);
+    umf_ba_global_free(memoryTarget, sizeof(umf_memory_target_t));
 }

src/memory_target.h

@@ -24,9 +24,6 @@ typedef struct umf_memory_target_ops_t umf_memory_target_ops_t;
 typedef struct umf_memory_target_t {
     const umf_memory_target_ops_t *ops;
     void *priv;
-
-    // saved pointer to the global base allocator
-    umf_ba_pool_t *base_allocator;
 } umf_memory_target_t;
 
 typedef umf_memory_target_t *umf_memory_target_handle_t;

src/memory_targets/memory_target_numa.c

@@ -22,9 +22,6 @@
 
 struct numa_memory_target_t {
     size_t id;
-
-    // saved pointer to the global base allocator
-    umf_ba_pool_t *base_allocator;
 };
 
 static umf_result_t numa_initialize(void *params, void **memTarget) {
@@ -35,27 +32,19 @@ static umf_result_t numa_initialize(void *params, void **memTarget) {
     struct umf_numa_memory_target_config_t *config =
         (struct umf_numa_memory_target_config_t *)params;
 
-    umf_ba_pool_t *base_allocator =
-        umf_ba_get_pool(sizeof(struct numa_memory_target_t));
-    if (!base_allocator) {
-        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
-    }
-
-    struct numa_memory_target_t *numaTarget = umf_ba_alloc(base_allocator);
+    struct numa_memory_target_t *numaTarget =
+        umf_ba_global_alloc(sizeof(struct numa_memory_target_t));
     if (!numaTarget) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
-    numaTarget->base_allocator = base_allocator;
     numaTarget->id = config->id;
     *memTarget = numaTarget;
     return UMF_RESULT_SUCCESS;
 }
 
 static void numa_finalize(void *memTarget) {
-    struct numa_memory_target_t *numaTarget =
-        (struct numa_memory_target_t *)memTarget;
-    umf_ba_free(numaTarget->base_allocator, memTarget);
+    umf_ba_global_free(memTarget, sizeof(struct numa_memory_target_t));
 }
 
 // sets maxnode and allocates and initializes mask based on provided memory targets

src/memspace.c

@@ -12,6 +12,7 @@
 
 #include <umf/memspace.h>
 
+#include "base_alloc_global.h"
 #include "memory_target.h"
 #include "memory_target_ops.h"
 #include "memspace_internal.h"
@@ -105,5 +106,5 @@ void umfMemspaceDestroy(umf_memspace_handle_t memspace) {
     }
 
     umf_ba_linear_destroy(memspace->linear_allocator);
-    umf_ba_free(memspace->base_allocator, memspace);
+    umf_ba_global_free(memspace, sizeof(struct umf_memspace_t));
 }

src/memspace_internal.h

@@ -24,9 +24,6 @@ struct umf_memspace_t {
     size_t size;
     umf_memory_target_handle_t *nodes;
 
-    // saved pointer to the global base allocator
-    umf_ba_pool_t *base_allocator;
-
     // own local linear base allocator
     umf_ba_linear_pool_t *linear_allocator;
 };