Global base alloc free function now does not require size param

src/base_alloc/base_alloc_global.c

@@ -14,13 +14,16 @@
 #include "base_alloc.h"
 #include "base_alloc_global.h"
 #include "base_alloc_internal.h"
+#include "utils_common.h"
 #include "utils_concurrency.h"
 #include "utils_math.h"
 
 // global base allocator used by all providers and pools
 static UTIL_ONCE_FLAG ba_is_initialized = UTIL_ONCE_FLAG_INIT;
 
-// allocation classes need to be powers of 2
+#define ALLOC_METADATA_SIZE (sizeof(size_t))
+
+// allocation classes need to be consecutive powers of 2
 #define ALLOCATION_CLASSES                                                     \
     { 16, 32, 64, 128 }
 #define NUM_ALLOCATION_CLASSES 4
@@ -82,41 +85,123 @@ static int size_to_idx(size_t size) {
     return index;
 }
 
-void *umf_ba_global_alloc(size_t size) {
+// stores metadata just before 'ptr' and returns beginning of usable
+// space to the user. metadata consists of 'size' that is the allocation
+// size and 'offset' that specifies how far is the returned ptr from
+// the origin ptr (used for aligned alloc)
+static void *add_metadata_and_align(void *ptr, size_t size, size_t alignment) {
+    assert(size < (1ULL << 32));
+    assert(alignment < (1ULL << 32));
+    assert(ptr);
+
+    void *user_ptr;
+    if (alignment <= ALLOC_METADATA_SIZE) {
+        user_ptr = (void *)((uintptr_t)ptr + ALLOC_METADATA_SIZE);
+    } else {
+        user_ptr =
+            (void *)ALIGN_UP((uintptr_t)ptr + ALLOC_METADATA_SIZE, alignment);
+    }
+
+    size_t ptr_offset_from_original = (uintptr_t)user_ptr - (uintptr_t)ptr;
+    assert(ptr_offset_from_original < (1ULL << 32));
+
+    size_t *metadata_loc = (size_t *)((char *)user_ptr - ALLOC_METADATA_SIZE);
+    *metadata_loc = size | (ptr_offset_from_original << 32);
+
+    return user_ptr;
+}
+
+// return original ptr (the one that has been passed to add_metadata_and_align)
+// along with total allocation size (needed to find proper alloc class
+// in free) and usable size
+static void *get_original_alloc(void *user_ptr, size_t *total_size,
+                                size_t *usable_size) {
+    assert(user_ptr);
+
+    size_t *metadata_loc = (size_t *)((char *)user_ptr - ALLOC_METADATA_SIZE);
+
+    size_t stored_size = *metadata_loc & ((1ULL << 32) - 1);
+    size_t ptr_offset_from_original = *metadata_loc >> 32;
+
+    void *original_ptr =
+        (void *)((uintptr_t)user_ptr - ptr_offset_from_original);
+
+    if (total_size) {
+        *total_size = stored_size;
+    }
+
+    if (usable_size) {
+        *usable_size = stored_size - ptr_offset_from_original;
+    }
+
+    return original_ptr;
+}
+
+void *umf_ba_global_aligned_alloc(size_t size, size_t alignment) {
     util_init_once(&ba_is_initialized, umf_ba_create_global);
 
+    // for metadata
+    size += ALLOC_METADATA_SIZE;
+
+    if (alignment > ALLOC_METADATA_SIZE) {
+        size += alignment;
+    }
+
     if (size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
 #ifndef NDEBUG
         fprintf(stderr,
                 "base_alloc: allocation size larger than the biggest "
                 "allocation class. Falling back to OS memory allocation.\n");
 #endif
-        return ba_os_alloc(size);
+        return add_metadata_and_align(ba_os_alloc(size), size, alignment);
     }
 
     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 add_metadata_and_align(ba_os_alloc(size), size, alignment);
     }
 
-    return umf_ba_alloc(BASE_ALLOC.ac[ac_index]);
+    return add_metadata_and_align(umf_ba_alloc(BASE_ALLOC.ac[ac_index]), size,
+                                  alignment);
 }
 
-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);
+void *umf_ba_global_alloc(size_t size) {
+    return umf_ba_global_aligned_alloc(size, ALLOC_METADATA_SIZE);
+}
+
+void umf_ba_global_free(void *ptr) {
+    if (!ptr) {
         return;
     }
 
-    int ac_index = size_to_idx(size);
+    size_t total_size;
+    ptr = get_original_alloc(ptr, &total_size, NULL);
+
+    if (total_size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
+        ba_os_free(ptr, total_size);
+        return;
+    }
+
+    int ac_index = size_to_idx(total_size);
     if (!BASE_ALLOC.ac[ac_index]) {
         // if creating ac failed, memory must have been allocated by os
-        ba_os_free(ptr, size);
+        ba_os_free(ptr, total_size);
         return;
     }
 
     umf_ba_free(BASE_ALLOC.ac[ac_index], ptr);
 }
+
+size_t umf_ba_global_malloc_usable_size(void *ptr) {
+    if (!ptr) {
+        return 0;
+    }
+
+    size_t usable_size;
+    get_original_alloc(ptr, NULL, &usable_size);
+
+    return usable_size;
+}

src/base_alloc/base_alloc_global.h

@@ -15,8 +15,10 @@ extern "C" {
 #endif
 
 void *umf_ba_global_alloc(size_t size);
-void umf_ba_global_free(void *ptr, size_t size);
+void umf_ba_global_free(void *ptr);
 void umf_ba_destroy_global(void);
+size_t umf_ba_global_malloc_usable_size(void *ptr);
+void *umf_ba_global_aligned_alloc(size_t size, size_t alignment);
 
 #ifdef __cplusplus
 }

src/critnib/critnib.c

@@ -211,7 +211,7 @@ struct critnib *critnib_new(void) {
 err_util_mutex_destroy:
     util_mutex_destroy_not_free(&c->mutex);
 err_free_critnib:
-    umf_ba_global_free(c, sizeof(struct critnib));
+    umf_ba_global_free(c);
     return NULL;
 }
 
@@ -261,7 +261,7 @@ void critnib_delete(struct critnib *c) {
 
     umf_ba_destroy(c->pool_nodes);
     umf_ba_destroy(c->pool_leaves);
-    umf_ba_global_free(c, sizeof(struct critnib));
+    umf_ba_global_free(c);
 }
 
 /*

src/memory_pool_default.c

@@ -39,7 +39,7 @@ umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
     pool->ops = *ops;
     ret = ops->initialize(pool->provider, params, &pool->pool_priv);
     if (ret != UMF_RESULT_SUCCESS) {
-        umf_ba_global_free(pool, sizeof(umf_memory_pool_t));
+        umf_ba_global_free(pool);
         return ret;
     }
 
@@ -56,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_global_free(hPool, sizeof(umf_memory_pool_t));
+    umf_ba_global_free(hPool);
 }
 
 umf_result_t umfFree(void *ptr) {

src/memory_pool_tracking.c

@@ -54,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_global_free(pool, sizeof(umf_memory_pool_t));
+    umf_ba_global_free(pool);
     return ret;
 }
 
@@ -69,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_global_free(hPool, sizeof(umf_memory_pool_t));
+    umf_ba_global_free(hPool);
 }
 
 umf_result_t umfFree(void *ptr) {

src/memory_provider.c

@@ -45,7 +45,7 @@ umf_result_t umfMemoryProviderCreate(const umf_memory_provider_ops_t *ops,
     void *provider_priv;
     umf_result_t ret = ops->initialize(params, &provider_priv);
     if (ret != UMF_RESULT_SUCCESS) {
-        umf_ba_global_free(provider, sizeof(umf_memory_provider_t));
+        umf_ba_global_free(provider);
         return ret;
     }
 
@@ -58,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_global_free(hProvider, sizeof(umf_memory_provider_t));
+    umf_ba_global_free(hProvider);
 }
 
 static void

src/memory_target.c

@@ -37,7 +37,7 @@ umf_result_t umfMemoryTargetCreate(const umf_memory_target_ops_t *ops,
     void *target_priv;
     umf_result_t ret = ops->initialize(params, &target_priv);
     if (ret != UMF_RESULT_SUCCESS) {
-        umf_ba_global_free(target, sizeof(umf_memory_target_t));
+        umf_ba_global_free(target);
         return ret;
     }
 
@@ -51,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_global_free(memoryTarget, sizeof(umf_memory_target_t));
+    umf_ba_global_free(memoryTarget);
 }

src/memory_targets/memory_target_numa.c

@@ -42,9 +42,7 @@ static umf_result_t numa_initialize(void *params, void **memTarget) {
     return UMF_RESULT_SUCCESS;
 }
 
-static void numa_finalize(void *memTarget) {
-    umf_ba_global_free(memTarget, sizeof(struct numa_memory_target_t));
-}
+static void numa_finalize(void *memTarget) { umf_ba_global_free(memTarget); }
 
 // sets maxnode and allocates and initializes mask based on provided memory targets
 static umf_result_t
@@ -93,7 +91,7 @@ numa_targets_create_nodemask(struct numa_memory_target_t **targets,
     hwloc_bitmap_free(bitmap);
 
     if (ret) {
-        umf_ba_global_free(nodemask, *mask_size);
+        umf_ba_global_free(nodemask);
         return UMF_RESULT_ERROR_UNKNOWN;
     }
 
@@ -131,7 +129,7 @@ static enum umf_result_t numa_memory_provider_create_from_memspace(
     umf_memory_provider_handle_t numaProvider = NULL;
     ret = umfMemoryProviderCreate(umfOsMemoryProviderOps(), &params,
                                   &numaProvider);
-    umf_ba_global_free(nodemask, nodemask_size);
+    umf_ba_global_free(nodemask);
     if (ret) {
         return ret;
     }

src/memspace.c

@@ -70,7 +70,7 @@ umf_result_t umfPoolCreateFromMemspace(umf_memspace_handle_t memspace,
     ret = memspace->nodes[0]->ops->pool_create_from_memspace(
         memspace, privs, memspace->size, policy, pool);
 
-    umf_ba_global_free(privs, sizeof(void *) * memspace->size);
+    umf_ba_global_free(privs);
 
     return ret;
 }
@@ -95,7 +95,7 @@ umfMemoryProviderCreateFromMemspace(umf_memspace_handle_t memspace,
     ret = memspace->nodes[0]->ops->memory_provider_create_from_memspace(
         memspace, privs, memspace->size, policy, provider);
 
-    umf_ba_global_free(privs, sizeof(void *) * memspace->size);
+    umf_ba_global_free(privs);
 
     return ret;
 }
@@ -106,7 +106,6 @@ void umfMemspaceDestroy(umf_memspace_handle_t memspace) {
         umfMemoryTargetDestroy(memspace->nodes[i]);
     }
 
-    umf_ba_global_free(memspace->nodes,
-                       memspace->size * sizeof(umf_memory_target_handle_t));
-    umf_ba_global_free(memspace, sizeof(struct umf_memspace_t));
+    umf_ba_global_free(memspace->nodes);
+    umf_ba_global_free(memspace);
 }

src/memspaces/memspace_host_all.c

@@ -62,7 +62,7 @@ static umf_result_t umfMemspaceHostAllCreate(umf_memspace_handle_t *hMemspace) {
     umf_ret =
         umfMemspaceCreateFromNumaArray(nodeIds, (size_t)nNodes, hMemspace);
 
-    umf_ba_global_free(nodeIds, nNodes * sizeof(size_t));
+    umf_ba_global_free(nodeIds);
 
 err_topology_destroy:
     hwloc_topology_destroy(topology);

src/memspaces/memspace_numa.c

@@ -52,12 +52,11 @@ umfMemspaceCreateFromNumaArray(size_t *nodeIds, size_t numIds,
     return UMF_RESULT_SUCCESS;
 
 err_target_create:
-    umf_ba_global_free(memspace->nodes,
-                       memspace->size * sizeof(umf_memory_target_handle_t));
+    umf_ba_global_free(memspace->nodes);
     for (size_t i = 0; i < nodeIdx; i++) {
         umfMemoryTargetDestroy(memspace->nodes[i]);
     }
 err_nodes_alloc:
-    umf_ba_global_free(memspace, sizeof(struct umf_memspace_t));
+    umf_ba_global_free(memspace);
     return ret;
 }

src/pool/pool_jemalloc.c

@@ -395,7 +395,7 @@ static umf_result_t je_initialize(umf_memory_provider_handle_t provider,
     return UMF_RESULT_SUCCESS;
 
 err_free_pool:
-    umf_ba_global_free(pool, sizeof(jemalloc_memory_pool_t));
+    umf_ba_global_free(pool);
     return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
 }
 
@@ -406,7 +406,7 @@ static void je_finalize(void *pool) {
     snprintf(cmd, sizeof(cmd), "arena.%u.destroy", je_pool->arena_index);
     mallctl(cmd, NULL, 0, NULL, 0);
     pool_by_arena_index[je_pool->arena_index] = NULL;
-    umf_ba_global_free(je_pool, sizeof(jemalloc_memory_pool_t));
+    umf_ba_global_free(je_pool);
 }
 
 static size_t je_malloc_usable_size(void *pool, void *ptr) {

src/pool/pool_scalable.c

@@ -170,7 +170,7 @@ static void tbb_pool_finalize(void *pool) {
     tbb_memory_pool_t *pool_data = (tbb_memory_pool_t *)pool;
     pool_data->tbb_callbacks.pool_destroy(pool_data->tbb_pool);
     dlclose(pool_data->tbb_callbacks.lib_handle);
-    umf_ba_global_free(pool_data, sizeof(tbb_memory_pool_t));
+    umf_ba_global_free(pool_data);
 }
 
 static void *tbb_malloc(void *pool, size_t size) {

src/provider/provider_os_memory.c

@@ -278,7 +278,7 @@ static umf_result_t os_initialize(void *params, void **provider) {
 err_destroy_hwloc_topology:
     hwloc_topology_destroy(os_provider->topo);
 err_free_os_provider:
-    umf_ba_global_free(os_provider, sizeof(os_memory_provider_t));
+    umf_ba_global_free(os_provider);
     return ret;
 }
 
@@ -291,7 +291,7 @@ static void os_finalize(void *provider) {
     os_memory_provider_t *os_provider = provider;
     hwloc_bitmap_free(os_provider->nodeset);
     hwloc_topology_destroy(os_provider->topo);
-    umf_ba_global_free(os_provider, sizeof(os_memory_provider_t));
+    umf_ba_global_free(os_provider);
 }
 
 static umf_result_t os_get_min_page_size(void *provider, void *ptr,