fix alloc / free aligned chunk address calc in disjoint pool

Author: bratpiorka

src/pool/pool_disjoint.c

@@ -10,6 +10,8 @@
 // Temporary solution for disabling memory poisoning. This is needed because
 // AddressSanitizer does not support memory poisoning for GPU allocations.
 // More info: https://github.com/oneapi-src/unified-memory-framework/issues/634
+
+// TODO - add a param to disjoint pool to disable memory poisoning
 #ifndef POISON_MEMORY
 #undef __SANITIZE_ADDRESS__
 #endif
@@ -94,9 +96,6 @@ static slab_t *create_slab(bucket_t *bucket) {
         goto free_slab_chunks;
     }
 
-    // TODO
-    // ASSERT_IS_ALIGNED((uintptr_t)slab->mem_ptr, bucket->size);
-
     // raw allocation is not available for user so mark it as inaccessible
     utils_annotate_memory_inaccessible(slab->mem_ptr, slab->slab_size);
 
@@ -178,7 +177,8 @@ static void slab_free_chunk(slab_t *slab, void *ptr) {
     // Even if the pointer p was previously aligned, it's still inside the
     // corresponding chunk, so we get the correct index here.
     size_t chunk_idx =
-        ((uintptr_t)ptr - (uintptr_t)slab->mem_ptr) / slab->bucket->size;
+        floor((double)((uintptr_t)ptr - (uintptr_t)slab->mem_ptr) /
+              slab->bucket->size);
 
     // Make sure that the chunk was allocated
     assert(slab->chunks[chunk_idx] && "double free detected");
@@ -738,8 +738,6 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
         }
     }
 
-    utils_mutex_unlock(&bucket->bucket_lock);
-
     if (disjoint_pool->params.pool_trace > 2) {
         LOG_DEBUG("Allocated %8zu %s bytes aligned at %zu from %s -> %p", size,
                   disjoint_pool->params.name, alignment,
@@ -749,6 +747,8 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
     void *aligned_ptr = (void *)ALIGN_UP_SAFE((size_t)ptr, alignment);
     VALGRIND_DO_MEMPOOL_ALLOC(disjoint_pool, aligned_ptr, size);
     utils_annotate_memory_undefined(aligned_ptr, size);
+
+    utils_mutex_unlock(&bucket->bucket_lock);
     return aligned_ptr;
 }
 
@@ -804,10 +804,20 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
 
     bucket_t *bucket = slab->bucket;
 
-    VALGRIND_DO_MEMPOOL_FREE(pool, ptr);
     utils_mutex_lock(&bucket->bucket_lock);
 
-    utils_annotate_memory_inaccessible(ptr, bucket->size);
+    // TODO valgrind
+    VALGRIND_DO_MEMPOOL_FREE(pool, ptr);
+
+    // Get the unalgined pointer
+    // NOTE: the base pointer slab->mem_ptr needn't to be aligned to bucket size
+    size_t chunk_idx =
+        floor((double)((uintptr_t)ptr - (uintptr_t)slab->mem_ptr) /
+              slab->bucket->size);
+    void *unaligned_ptr =
+        (void *)((uintptr_t)slab->mem_ptr + chunk_idx * slab->bucket->size);
+
+    utils_annotate_memory_inaccessible(unaligned_ptr, bucket->size);
     bucket_free_chunk(bucket, ptr, slab, &to_pool);
 
     if (disjoint_pool->params.pool_trace > 1) {

src/pool/pool_disjoint_internal.h

@@ -11,6 +11,7 @@
 #include <assert.h>
 #include <ctype.h>
 #include <errno.h>
+#include <math.h>
 #include <stdbool.h>
 #include <stdlib.h>
 #include <string.h>