Implement umfMemoryProviderAllocation[Split/Merge]

include/umf/memory_provider.h

@@ -156,6 +156,34 @@ umfMemoryProviderGetName(umf_memory_provider_handle_t hProvider);
 ///
 UMF_EXPORT umf_memory_provider_handle_t umfGetLastFailedMemoryProvider(void);
 
+///
+/// @brief Splits a coarse grain allocation into 2 adjacent allocations that
+///        can be managed (freed) separately.
+/// @param hProvider handle to the memory provider
+/// @param ptr pointer to the beginning of the allocation
+/// @param totalSize total size of the allocation to be split
+/// @param firstSize size of the first new allocation, second allocation
+//         has a size equal to totalSize - firstSize
+/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure
+///
+UMF_EXPORT umf_result_t
+umfMemoryProviderAllocationSplit(umf_memory_provider_handle_t hProvider,
+                                 void *ptr, size_t totalSize, size_t firstSize);
+
+///
+/// @brief Merges two coarse grain allocations into a single allocation that
+///        can be managed (freed) as a whole.
+/// @param hProvider handle to the memory provider
+/// @param lowPtr pointer to the first allocation
+/// @param highPtr pointer to the second allocation (should be > lowPtr)
+/// @param totalSize size of a new merged allocation. Should be equal
+///        to the sum of sizes of allocations beginning at lowPtr and highPtr
+/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure
+///
+UMF_EXPORT umf_result_t
+umfMemoryProviderAllocationMerge(umf_memory_provider_handle_t hProvider,
+                                 void *lowPtr, void *highPtr, size_t totalSize);
+
 #ifdef __cplusplus
 }
 #endif

include/umf/memory_provider_ops.h

@@ -139,6 +139,32 @@ typedef struct umf_memory_provider_ops_t {
     /// @return pointer to a string containing the name of the \p provider
     ///
     const char *(*get_name)(void *provider);
+
+    ///
+    /// @brief Splits a coarse grain allocation into 2 adjacent allocations that
+    ///        can be managed (freed) separately.
+    /// @param hProvider handle to the memory provider
+    /// @param ptr pointer to the beginning of the allocation
+    /// @param totalSize total size of the allocation to be split
+    /// @param firstSize size of the first new allocation, second allocation
+    //         has a size equal to totalSize - firstSize
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure
+    ///
+    umf_result_t (*allocation_split)(void *hProvider, void *ptr,
+                                     size_t totalSize, size_t firstSize);
+
+    ///
+    /// @brief Merges two coarse grain allocations into a single allocation that
+    ///        can be managed (freed) as a whole.
+    /// @param hProvider handle to the memory provider
+    /// @param lowPtr pointer to the first allocation
+    /// @param highPtr pointer to the second allocation (should be > lowPtr)
+    /// @param totalSize size of a new merged allocation. Should be equal
+    ///        to the sum of sizes of allocations beginning at lowPtr and highPtr
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure
+    ///
+    umf_result_t (*allocation_merge)(void *hProvider, void *lowPtr,
+                                     void *highPtr, size_t totalSize);
 } umf_memory_provider_ops_t;
 
 #ifdef __cplusplus

src/memory_provider.c

@@ -11,6 +11,7 @@
 #include <umf/memory_provider.h>
 
 #include <assert.h>
+#include <stdio.h>
 #include <stdlib.h>
 
 typedef struct umf_memory_provider_t {
@@ -129,3 +130,46 @@ const char *umfMemoryProviderGetName(umf_memory_provider_handle_t hProvider) {
 umf_memory_provider_handle_t umfGetLastFailedMemoryProvider(void) {
     return *umfGetLastFailedMemoryProviderPtr();
 }
+
+umf_result_t
+umfMemoryProviderAllocationSplit(umf_memory_provider_handle_t hProvider,
+                                 void *ptr, size_t totalSize,
+                                 size_t firstSize) {
+    if (!ptr) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (firstSize == 0 || totalSize == 0) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (firstSize >= totalSize) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t res = hProvider->ops.allocation_split(
+        hProvider->provider_priv, ptr, totalSize, firstSize);
+    checkErrorAndSetLastProvider(res, hProvider);
+    return res;
+}
+
+umf_result_t
+umfMemoryProviderAllocationMerge(umf_memory_provider_handle_t hProvider,
+                                 void *lowPtr, void *highPtr,
+                                 size_t totalSize) {
+    if (!lowPtr || !highPtr) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (totalSize == 0) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if ((uintptr_t)lowPtr >= (uintptr_t)highPtr) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if ((uintptr_t)highPtr - (uintptr_t)lowPtr > totalSize) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t res = hProvider->ops.allocation_merge(
+        hProvider->provider_priv, lowPtr, highPtr, totalSize);
+    checkErrorAndSetLastProvider(res, hProvider);
+    return res;
+}

src/pool/pool_jemalloc.c

@@ -210,15 +210,15 @@ static bool arena_extent_split(extent_hooks_t *extent_hooks, void *addr,
                                size_t size, size_t size_a, size_t size_b,
                                bool committed, unsigned arena_ind) {
     (void)extent_hooks; // unused
-    (void)addr;         // unused
-    (void)size;         // unused
-    (void)size_a;       // unused
-    (void)size_b;       // unused
     (void)committed;    // unused
-    (void)arena_ind;    // unused
+    (void)size_b;
 
-    // TODO: add this function to the provider API to support Windows and USM
-    return true; // true means failure (unsupported)
+    assert(size_a + size_b == size);
+
+    jemalloc_memory_pool_t *pool = get_pool_by_arena_index(arena_ind);
+    assert(pool);
+    return umfMemoryProviderAllocationSplit(pool->provider, addr, size,
+                                            size_a) != UMF_RESULT_SUCCESS;
 }
 
 // arena_extent_merge - an extent merge function conforms to the extent_merge_t type and optionally
@@ -231,15 +231,13 @@ static bool arena_extent_merge(extent_hooks_t *extent_hooks, void *addr_a,
                                size_t size_a, void *addr_b, size_t size_b,
                                bool committed, unsigned arena_ind) {
     (void)extent_hooks; // unused
-    (void)addr_a;       // unused
-    (void)addr_b;       // unused
-    (void)size_a;       // unused
-    (void)size_b;       // unused
     (void)committed;    // unused
-    (void)arena_ind;    // unused
 
-    // TODO: add this function to the provider API to support Windows and USM
-    return true; // true means failure (unsupported)
+    jemalloc_memory_pool_t *pool = get_pool_by_arena_index(arena_ind);
+    assert(pool);
+    return umfMemoryProviderAllocationMerge(pool->provider, addr_a, addr_b,
+                                            size_a + size_b) !=
+           UMF_RESULT_SUCCESS;
 }
 
 // The extent_hooks_t structure comprises function pointers which are described individually below.

src/pool/pool_scalable.c

@@ -121,7 +121,7 @@ static void tbb_raw_free_wrapper(intptr_t pool_id, void *ptr, size_t bytes) {
         TLS_last_free_error = ret;
         fprintf(
             stderr,
-            "Memory provider failed to free memory, addr = %p, size = %lu\n",
+            "Memory provider failed to free memory, addr = %p, size = %zu\n",
             ptr, bytes);
     }
 }

src/provider/provider_os_memory.c

@@ -449,6 +449,26 @@ static const char *os_get_name(void *provider) {
     return "OS";
 }
 
+static umf_result_t os_allocation_split(void *provider, void *ptr,
+                                        size_t totalSize, size_t firstSize) {
+    (void)provider;
+    (void)ptr;
+    (void)totalSize;
+    (void)firstSize;
+    // nop
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t os_allocation_merge(void *provider, void *lowPtr,
+                                        void *highPtr, size_t totalSize) {
+    (void)provider;
+    (void)lowPtr;
+    (void)highPtr;
+    (void)totalSize;
+    // nop
+    return UMF_RESULT_SUCCESS;
+}
+
 umf_memory_provider_ops_t UMF_OS_MEMORY_PROVIDER_OPS = {
     .version = UMF_VERSION_CURRENT,
     .initialize = os_initialize,
@@ -461,4 +481,5 @@ umf_memory_provider_ops_t UMF_OS_MEMORY_PROVIDER_OPS = {
     .purge_lazy = os_purge_lazy,
     .purge_force = os_purge_force,
     .get_name = os_get_name,
-};
+    .allocation_split = os_allocation_split,
+    .allocation_merge = os_allocation_merge};