a

bratpiorka · bratpiorka · commit 8f92d311445e · 2025-02-24T15:47:51.000Z
diff --git a/src/critnib/critnib.c b/src/critnib/critnib.c
@@ -1,6 +1,6 @@
 /*
  *
- * Copyright (C) 2023-2024 Intel Corporation
+ * Copyright (C) 2023-2025 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -133,24 +133,6 @@ struct critnib {
     struct utils_mutex_t mutex; /* writes/removes */
 };
 
-/*
- * atomic load
- */
-static void load(void *src, void *dst) {
-    utils_atomic_load_acquire((word *)src, (word *)dst);
-}
-
-static void load64(uint64_t *src, uint64_t *dst) {
-    utils_atomic_load_acquire(src, dst);
-}
-
-/*
- * atomic store
- */
-static void store(void *dst, void *src) {
-    utils_atomic_store_release((word *)dst, (word)src);
-}
-
 /*
  * internal: is_leaf -- check tagged pointer for leafness
  */
@@ -343,10 +325,8 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
 
     struct critnib_node *n = c->root;
     if (!n) {
-        store(&c->root, kn);
-
+        utils_atomic_store_release_ptr(&c->root, kn);
         utils_mutex_unlock(&c->mutex);
-
         return 0;
     }
 
@@ -361,7 +341,8 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
 
     if (!n) {
         n = prev;
-        store(&n->child[slice_index(key, n->shift)], kn);
+        utils_atomic_store_release_ptr(&n->child[slice_index(key, n->shift)],
+                                       kn);
 
         utils_mutex_unlock(&c->mutex);
 
@@ -406,7 +387,7 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
     m->child[slice_index(path, sh)] = n;
     m->shift = sh;
     m->path = key & path_mask(sh);
-    store(parent, m);
+    utils_atomic_store_release_ptr(parent, m);
 
     utils_mutex_unlock(&c->mutex);
 
@@ -427,7 +408,8 @@ void *critnib_remove(struct critnib *c, word key) {
         goto not_found;
     }
 
-    word del = (utils_atomic_increment(&c->remove_count) - 1) % DELETED_LIFE;
+    word del =
+        (utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
     free_node(c, c->pending_del_nodes[del]);
     free_leaf(c, c->pending_del_leaves[del]);
     c->pending_del_nodes[del] = NULL;
@@ -436,7 +418,7 @@ void *critnib_remove(struct critnib *c, word key) {
     if (is_leaf(n)) {
         k = to_leaf(n);
         if (k->key == key) {
-            store(&c->root, NULL);
+            utils_atomic_store_release_ptr(&c->root, NULL);
             goto del_leaf;
         }
 
@@ -466,7 +448,7 @@ void *critnib_remove(struct critnib *c, word key) {
         goto not_found;
     }
 
-    store(&n->child[slice_index(key, n->shift)], NULL);
+    utils_atomic_store_release_ptr(&n->child[slice_index(key, n->shift)], NULL);
 
     /* Remove the node if there's only one remaining child. */
     int ochild = -1;
@@ -482,7 +464,7 @@ void *critnib_remove(struct critnib *c, word key) {
 
     ASSERTne(ochild, -1);
 
-    store(n_parent, n->child[ochild]);
+    utils_atomic_store_release_ptr(n_parent, n->child[ochild]);
     c->pending_del_nodes[del] = n;
 
 del_leaf:
@@ -511,22 +493,23 @@ void *critnib_get(struct critnib *c, word key) {
     do {
         struct critnib_node *n;
 
-        load64(&c->remove_count, &wrs1);
-        load(&c->root, &n);
+        utils_atomic_load_acquire_u64(&c->remove_count, &wrs1);
+        utils_atomic_load_acquire_ptr(&c->root, (void **)&n);
 
         /*
 		 * critbit algorithm: dive into the tree, looking at nothing but
 		 * each node's critical bit^H^H^Hnibble.  This means we risk
 		 * going wrong way if our path is missing, but that's ok...
 		 */
         while (n && !is_leaf(n)) {
-            load(&n->child[slice_index(key, n->shift)], &n);
+            utils_atomic_load_acquire_ptr(&n->child[slice_index(key, n->shift)],
+                                          (void **)&n);
         }
 
         /* ... as we check it at the end. */
         struct critnib_leaf *k = to_leaf(n);
         res = (n && k->key == key) ? k->value : NULL;
-        load64(&c->remove_count, &wrs2);
+        utils_atomic_load_acquire_u64(&c->remove_count, &wrs2);
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
     return res;
@@ -597,7 +580,7 @@ static struct critnib_leaf *find_le(struct critnib_node *__restrict n,
     /* recursive call: follow the path */
     {
         struct critnib_node *m;
-        load(&n->child[nib], &m);
+        utils_atomic_load_acquire_ptr(&n->child[nib], (void **)&m);
         struct critnib_leaf *k = find_le(m, key);
         if (k) {
             return k;
@@ -611,7 +594,7 @@ static struct critnib_leaf *find_le(struct critnib_node *__restrict n,
 	 */
     for (; nib > 0; nib--) {
         struct critnib_node *m;
-        load(&n->child[nib - 1], &m);
+        utils_atomic_load_acquire_ptr(&n->child[nib - 1], (void **)&m);
         if (m) {
             n = m;
             if (is_leaf(n)) {
@@ -635,12 +618,12 @@ void *critnib_find_le(struct critnib *c, word key) {
     void *res;
 
     do {
-        load64(&c->remove_count, &wrs1);
+        utils_atomic_load_acquire_u64(&c->remove_count, &wrs1);
         struct critnib_node *n; /* avoid a subtle TOCTOU */
-        load(&c->root, &n);
+        utils_atomic_load_acquire_ptr(&c->root, (void **)&n);
         struct critnib_leaf *k = n ? find_le(n, key) : NULL;
         res = k ? k->value : NULL;
-        load64(&c->remove_count, &wrs2);
+        utils_atomic_load_acquire_u64(&c->remove_count, &wrs2);
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
     return res;
@@ -694,7 +677,7 @@ static struct critnib_leaf *find_ge(struct critnib_node *__restrict n,
     unsigned nib = slice_index(key, n->shift);
     {
         struct critnib_node *m;
-        load(&n->child[nib], &m);
+        utils_atomic_load_acquire_ptr(&n->child[nib], (void **)&m);
         struct critnib_leaf *k = find_ge(m, key);
         if (k) {
             return k;
@@ -703,7 +686,7 @@ static struct critnib_leaf *find_ge(struct critnib_node *__restrict n,
 
     for (; nib < NIB; nib++) {
         struct critnib_node *m;
-        load(&n->child[nib + 1], &m);
+        utils_atomic_load_acquire_ptr(&n->child[nib + 1], (void **)&m);
         if (m) {
             n = m;
             if (is_leaf(n)) {
@@ -741,17 +724,18 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
     }
 
     do {
-        load64(&c->remove_count, &wrs1);
+        utils_atomic_load_acquire_u64(&c->remove_count, &wrs1);
         struct critnib_node *n;
-        load(&c->root, &n);
+        utils_atomic_load_acquire_ptr(&c->root, (void **)&n);
 
         if (dir < 0) {
             k = find_le(n, key);
         } else if (dir > 0) {
             k = find_ge(n, key);
         } else {
             while (n && !is_leaf(n)) {
-                load(&n->child[slice_index(key, n->shift)], &n);
+                utils_atomic_load_acquire_ptr(
+                    &n->child[slice_index(key, n->shift)], (void **)&n);
             }
 
             struct critnib_leaf *kk = to_leaf(n);
@@ -761,7 +745,7 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
             _rkey = k->key;
             _rvalue = k->value;
         }
-        load64(&c->remove_count, &wrs2);
+        utils_atomic_load_acquire_u64(&c->remove_count, &wrs2);
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
     if (k) {
diff --git a/src/ipc_cache.c b/src/ipc_cache.c
@@ -232,7 +232,7 @@ umf_result_t umfIpcOpenedCacheGet(ipc_opened_cache_handle_t cache,
 
 exit:
     if (ret == UMF_RESULT_SUCCESS) {
-        utils_atomic_increment(&entry->ref_count);
+        utils_atomic_increment_u64(&entry->ref_count);
         *retEntry = &entry->value;
     }
 
diff --git a/src/libumf.c b/src/libumf.c
@@ -24,10 +24,10 @@
 
 umf_memory_tracker_handle_t TRACKER = NULL;
 
-static unsigned long long umfRefCount = 0;
+static uint64_t umfRefCount = 0;
 
 int umfInit(void) {
-    if (utils_fetch_and_add64(&umfRefCount, 1) == 0) {
+    if (utils_fetch_and_add_u64(&umfRefCount, 1) == 0) {
         utils_log_init();
         TRACKER = umfMemoryTrackerCreate();
         if (!TRACKER) {
@@ -54,7 +54,7 @@ int umfInit(void) {
 }
 
 void umfTearDown(void) {
-    if (utils_fetch_and_add64(&umfRefCount, -1) == 1) {
+    if (utils_fetch_and_add_u64(&umfRefCount, -1) == 1) {
 #if !defined(_WIN32) && !defined(UMF_NO_HWLOC)
         umfMemspaceHostAllDestroy();
         umfMemspaceHighestCapacityDestroy();
diff --git a/src/pool/pool_disjoint.c b/src/pool/pool_disjoint.c
@@ -20,6 +20,7 @@
 #include "provider/provider_tracking.h"
 #include "uthash/utlist.h"
 #include "utils_common.h"
+#include "utils_concurrency.h"
 #include "utils_log.h"
 #include "utils_math.h"
 
@@ -34,7 +35,6 @@
 // Forward declarations
 static void bucket_update_stats(bucket_t *bucket, int in_use, int in_pool);
 static bool bucket_can_pool(bucket_t *bucket);
-static void bucket_decrement_pool(bucket_t *bucket);
 static slab_list_item_t *bucket_get_avail_slab(bucket_t *bucket,
                                                bool *from_pool);
 
@@ -316,6 +316,7 @@ static void bucket_free_chunk(bucket_t *bucket, void *ptr, slab_t *slab,
             assert(slab_it->val != NULL);
             pool_unregister_slab(bucket->pool, slab_it->val);
             DL_DELETE(bucket->available_slabs, slab_it);
+            assert(bucket->available_slabs_num > 0);
             bucket->available_slabs_num--;
             destroy_slab(slab_it->val);
         }
@@ -381,10 +382,17 @@ static slab_list_item_t *bucket_get_avail_slab(bucket_t *bucket,
         // Allocation from existing slab is treated as from pool for statistics.
         *from_pool = true;
         if (slab->num_chunks_allocated == 0) {
+            assert(bucket->chunked_slabs_in_pool > 0);
+            assert(bucket->shared_limits->total_size >=
+                   bucket_slab_alloc_size(bucket));
             // If this was an empty slab, it was in the pool.
             // Now it is no longer in the pool, so update count.
             --bucket->chunked_slabs_in_pool;
-            bucket_decrement_pool(bucket);
+            CACHE_ALIGNED size_t size_to_add =
+                -(long long)bucket_slab_alloc_size(bucket);
+            utils_fetch_and_add_u64(&bucket->shared_limits->total_size,
+                                    size_to_add);
+            bucket_update_stats(bucket, 1, -1);
         }
     }
 
@@ -420,36 +428,26 @@ static void bucket_update_stats(bucket_t *bucket, int in_use, int in_pool) {
         in_pool * bucket_slab_alloc_size(bucket);
 }
 
-static void bucket_decrement_pool(bucket_t *bucket) {
-    bucket_update_stats(bucket, 1, -1);
-    utils_fetch_and_add64(&bucket->shared_limits->total_size,
-                          -(long long)bucket_slab_alloc_size(bucket));
-}
-
 static bool bucket_can_pool(bucket_t *bucket) {
     size_t new_free_slabs_in_bucket;
 
     new_free_slabs_in_bucket = bucket->chunked_slabs_in_pool + 1;
 
     // we keep at most params.capacity slabs in the pool
     if (bucket_max_pooled_slabs(bucket) >= new_free_slabs_in_bucket) {
-        size_t pool_size = 0;
-        utils_atomic_load_acquire(&bucket->shared_limits->total_size,
-                                  &pool_size);
-        while (true) {
-            size_t new_pool_size = pool_size + bucket_slab_alloc_size(bucket);
-
-            if (bucket->shared_limits->max_size < new_pool_size) {
-                break;
-            }
-
-            if (utils_compare_exchange(&bucket->shared_limits->total_size,
-                                       &pool_size, &new_pool_size)) {
-                ++bucket->chunked_slabs_in_pool;
-
-                bucket_update_stats(bucket, -1, 1);
-                return true;
-            }
+
+        CACHE_ALIGNED size_t size_to_add = bucket_slab_alloc_size(bucket);
+        size_t previous_size = utils_fetch_and_add_u64(
+            &bucket->shared_limits->total_size, size_to_add);
+
+        if (previous_size + size_to_add <= bucket->shared_limits->max_size) {
+
+            ++bucket->chunked_slabs_in_pool;
+            bucket_update_stats(bucket, -1, 1);
+            return true;
+        } else {
+            utils_fetch_and_add_u64(&bucket->shared_limits->total_size,
+                                    -size_to_add);
         }
     }
 
@@ -895,7 +893,8 @@ umf_memory_pool_ops_t *umfDisjointPoolOps(void) {
 
 umf_disjoint_pool_shared_limits_t *
 umfDisjointPoolSharedLimitsCreate(size_t max_size) {
-    umf_disjoint_pool_shared_limits_t *ptr = umf_ba_global_alloc(sizeof(*ptr));
+    umf_disjoint_pool_shared_limits_t *ptr =
+        umf_ba_global_aligned_alloc(sizeof(*ptr), 64);
     if (ptr == NULL) {
         LOG_ERR("cannot allocate memory for disjoint pool shared limits");
         return NULL;
diff --git a/src/pool/pool_disjoint_internal.h b/src/pool/pool_disjoint_internal.h
@@ -101,8 +101,8 @@ typedef struct slab_t {
 } slab_t;
 
 typedef struct umf_disjoint_pool_shared_limits_t {
+    CACHE_ALIGNED size_t total_size; // requires atomic access
     size_t max_size;
-    size_t total_size; // requires atomic access
 } umf_disjoint_pool_shared_limits_t;
 
 typedef struct umf_disjoint_pool_params_t {
diff --git a/src/provider/provider_os_memory.c b/src/provider/provider_os_memory.c
@@ -934,7 +934,7 @@ static membind_t membindFirst(os_memory_provider_t *provider, void *addr,
 
     if (provider->mode == UMF_NUMA_MODE_INTERLEAVE) {
         assert(provider->part_size != 0);
-        size_t s = utils_fetch_and_add64(&provider->alloc_sum, size);
+        size_t s = utils_fetch_and_add_u64(&provider->alloc_sum, size);
         membind.node = (s / provider->part_size) % provider->nodeset_len;
         membind.bitmap = provider->nodeset[membind.node];
         membind.bind_size = ALIGN_UP(provider->part_size, membind.page_size);
diff --git a/src/provider/provider_tracking.c b/src/provider/provider_tracking.c
@@ -565,7 +565,7 @@ static umf_result_t trackingGetIpcHandle(void *provider, const void *ptr,
                 return ret;
             }
 
-            cache_value->handle_id = utils_atomic_increment(&IPC_HANDLE_ID);
+            cache_value->handle_id = utils_atomic_increment_u64(&IPC_HANDLE_ID);
             cache_value->ipcDataSize = ipcDataSize;
 
             int insRes = critnib_insert(p->ipcCache, (uintptr_t)ptr,
@@ -703,18 +703,20 @@ static umf_result_t trackingOpenIpcHandle(void *provider, void *providerIpcData,
     assert(cache_entry != NULL);
 
     void *mapped_ptr = NULL;
-    utils_atomic_load_acquire(&(cache_entry->mapped_base_ptr), &mapped_ptr);
+    utils_atomic_load_acquire_ptr(&(cache_entry->mapped_base_ptr),
+                                  (void **)&mapped_ptr);
     if (mapped_ptr == NULL) {
         utils_mutex_lock(&(cache_entry->mmap_lock));
-        utils_atomic_load_acquire(&(cache_entry->mapped_base_ptr), &mapped_ptr);
+        utils_atomic_load_acquire_ptr(&(cache_entry->mapped_base_ptr),
+                                      (void **)&mapped_ptr);
         if (mapped_ptr == NULL) {
             ret = upstreamOpenIPCHandle(p, providerIpcData,
                                         ipcUmfData->baseSize, &mapped_ptr);
             if (ret == UMF_RESULT_SUCCESS) {
                 // Put to the cache
                 cache_entry->mapped_size = ipcUmfData->baseSize;
-                utils_atomic_store_release(&(cache_entry->mapped_base_ptr),
-                                           mapped_ptr);
+                utils_atomic_store_release_ptr(&(cache_entry->mapped_base_ptr),
+                                               mapped_ptr);
             }
         }
         utils_mutex_unlock(&(cache_entry->mmap_lock));
diff --git a/src/utils/utils_concurrency.h b/src/utils/utils_concurrency.h
diff --git a/src/utils/utils_posix_concurrency.c b/src/utils/utils_posix_concurrency.c
diff --git a/src/utils/utils_windows_concurrency.c b/src/utils/utils_windows_concurrency.c

Original file line number	Diff line number	Diff line change
`@@ -232,7 +232,7 @@ umf_result_t umfIpcOpenedCacheGet(ipc_opened_cache_handle_t cache,`
`232`	`232`
`233`	`233`	`exit:`
`234`	`234`	`if (ret == UMF_RESULT_SUCCESS) {`
`235`		`- utils_atomic_increment(&entry->ref_count);`
	`235`	`+ utils_atomic_increment_u64(&entry->ref_count);`
`236`	`236`	`*retEntry = &entry->value;`
`237`	`237`	`}`
`238`	`238`