Lukasz and Rafal comments applied

Author: lslusarczyk

src/libumf.def

@@ -19,7 +19,6 @@ EXPORTS
     umfCUDAMemoryProviderParamsSetContext
     umfCUDAMemoryProviderParamsSetDevice
     umfCUDAMemoryProviderParamsSetMemoryType
-    umfLevelZeroMemoryProviderParamsSetResidentDevices
     umfDevDaxMemoryProviderOps
     umfDevDaxMemoryProviderParamsCreate
     umfDevDaxMemoryProviderParamsDestroy
@@ -42,6 +41,7 @@ EXPORTS
     umfLevelZeroMemoryProviderParamsSetContext
     umfLevelZeroMemoryProviderParamsSetDevice
     umfLevelZeroMemoryProviderParamsSetMemoryType
+    umfLevelZeroMemoryProviderParamsSetResidentDevices
     umfMemoryProviderAlloc
     umfMemoryProviderAllocationMerge
     umfMemoryProviderAllocationSplit

src/provider/provider_level_zero.c

@@ -422,6 +422,43 @@ static ze_relaxed_allocation_limits_exp_desc_t relaxed_device_allocation_desc =
      .pNext = NULL,
      .flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE};
 
+static umf_result_t ze_memory_provider_free_helper(void *provider, void *ptr,
+                                                   size_t bytes,
+                                                   int update_stats) {
+    if (ptr == NULL) {
+        return UMF_RESULT_SUCCESS;
+    }
+
+    ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
+    umf_result_t ret;
+    if (ze_provider->freePolicyFlags == 0) {
+        ret = ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
+    } else {
+        ze_memory_free_ext_desc_t desc = {
+            .stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
+            .pNext = NULL,
+            .freePolicy = ze_provider->freePolicyFlags};
+
+        ret = ze2umf_result(
+            g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
+    }
+
+    if (ret != UMF_RESULT_SUCCESS) {
+        return ret;
+    }
+
+    if (update_stats) {
+        provider_ctl_stats_free(ze_provider, bytes);
+    }
+
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
+                                            size_t bytes) {
+    return ze_memory_provider_free_helper(provider, ptr, bytes, 1);
+}
+
 static umf_result_t ze_memory_provider_alloc_helper(void *provider, size_t size,
                                                     size_t alignment,
                                                     int update_stats,
@@ -488,10 +525,17 @@ static umf_result_t ze_memory_provider_alloc_helper(void *provider, size_t size,
             *resultPtr, size);
         if (ze_result != ZE_RESULT_SUCCESS) {
             utils_read_unlock(&ze_provider->resident_device_rwlock);
-            LOG_DEBUG("making resident allocation %p of size:%lu on device %p "
-                      "failed with %d",
-                      *resultPtr, size, ze_provider->resident_device_handles[i],
-                      ze_result);
+            LOG_ERR("making resident allocation %p of size:%lu on device %p "
+                    "failed with 0x%x",
+                    *resultPtr, size, ze_provider->resident_device_handles[i],
+                    ze_result);
+            umf_result_t free_result =
+                ze_memory_provider_free(ze_provider, *resultPtr, size);
+            if (free_result != UMF_RESULT_SUCCESS) {
+                LOG_ERR("failed to free memory with: 0x%x after failed making "
+                        "resident, free fail ignored",
+                        free_result);
+            }
             return ze2umf_result(ze_result);
         }
         LOG_DEBUG("allocation %p of size:%lu made resident on device %p",
@@ -512,43 +556,6 @@ static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,
                                            resultPtr);
 }
 
-static umf_result_t ze_memory_provider_free_helper(void *provider, void *ptr,
-                                                   size_t bytes,
-                                                   int update_stats) {
-    if (ptr == NULL) {
-        return UMF_RESULT_SUCCESS;
-    }
-
-    ze_memory_provider_t *ze_provider = (ze_memory_provider_t *)provider;
-    umf_result_t ret;
-    if (ze_provider->freePolicyFlags == 0) {
-        ret = ze2umf_result(g_ze_ops.zeMemFree(ze_provider->context, ptr));
-    } else {
-        ze_memory_free_ext_desc_t desc = {
-            .stype = ZE_STRUCTURE_TYPE_MEMORY_FREE_EXT_DESC,
-            .pNext = NULL,
-            .freePolicy = ze_provider->freePolicyFlags};
-
-        ret = ze2umf_result(
-            g_ze_ops.zeMemFreeExt(ze_provider->context, &desc, ptr));
-    }
-
-    if (ret != UMF_RESULT_SUCCESS) {
-        return ret;
-    }
-
-    if (update_stats) {
-        provider_ctl_stats_free(ze_provider, bytes);
-    }
-
-    return UMF_RESULT_SUCCESS;
-}
-
-static umf_result_t ze_memory_provider_free(void *provider, void *ptr,
-                                            size_t bytes) {
-    return ze_memory_provider_free_helper(provider, ptr, bytes, 1);
-}
-
 static umf_result_t query_min_page_size(ze_memory_provider_t *ze_provider,
                                         size_t *min_page_size) {
     assert(min_page_size);
@@ -997,8 +1004,8 @@ static int ze_memory_provider_resident_device_change_helper(uintptr_t key,
 
     // TODO: add assertions to UMF and change it to be an assertion
     if (info->props.base != (void *)key) {
-        LOG_ERR("key:%p is different than base:%p", (void *)key,
-                info->props.base);
+        LOG_FATAL("key:%p is different than base:%p", (void *)key,
+                  info->props.base);
         abort();
     }
 
@@ -1060,8 +1067,7 @@ static umf_result_t ze_memory_provider_resident_device_change(void *provider,
         if (ze_provider->resident_device_count ==
             ze_provider->resident_device_capacity) {
             const uint32_t new_capacity =
-                ze_provider->resident_device_capacity * 2 +
-                1; // +1 to work also with old capacity == 0
+                ze_provider->resident_device_capacity + 1;
             ze_device_handle_t *new_handles =
                 umf_ba_global_alloc(sizeof(ze_device_handle_t) * new_capacity);
             if (new_handles == NULL) {
@@ -1087,7 +1093,6 @@ static umf_result_t ze_memory_provider_resident_device_change(void *provider,
         // found
         if (is_adding) {
             utils_write_unlock(&ze_provider->resident_device_rwlock);
-            // impossible for UR, should be an assertion
             LOG_ERR("trying to add resident device:%p but the device is "
                     "already a peer of provider:%p",
                     ze_device, provider);
@@ -1109,6 +1114,9 @@ static umf_result_t ze_memory_provider_resident_device_change(void *provider,
         .failed_changes = 0,
     };
 
+    // This is "hacky" and it will not work if someone uses pool without tracker
+    // or just use provider without pool. It can be solved by keeping track of
+    // allocations by the provider like in os_provider.
     umf_result_t result = umfMemoryTrackerIterateAll(
         &ze_memory_provider_resident_device_change_helper, &privData);
     if (result != UMF_RESULT_SUCCESS) {

test/common/CMakeLists.txt

@@ -1,4 +1,4 @@
-# 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
 
@@ -21,16 +21,18 @@ add_umf_library(
 target_include_directories(umf_test_common
                            PRIVATE ${UMF_CMAKE_SOURCE_DIR}/include)
 
-add_library(umf_test_mocks STATIC level_zero_mocks.cpp)
+if(UMF_LEVEL_ZERO_ENABLED)
+    add_library(umf_test_mocks STATIC level_zero_mocks.cpp)
 
-target_link_libraries(umf_test_mocks GTest::gmock)
+    target_link_libraries(umf_test_mocks GTest::gmock)
 
-target_include_directories(
-    umf_test_mocks
-    PRIVATE ${UMF_CMAKE_SOURCE_DIR}/include ${UMF_CMAKE_SOURCE_DIR}/src
-            ${UMF_CMAKE_SOURCE_DIR}/src/utils)
+    target_include_directories(
+        umf_test_mocks
+        PRIVATE ${UMF_CMAKE_SOURCE_DIR}/include ${UMF_CMAKE_SOURCE_DIR}/src
+                ${UMF_CMAKE_SOURCE_DIR}/src/utils)
 
-add_library(umf_ze_loopback SHARED ze_loopback.h ze_loopback.cpp)
-target_include_directories(
-    umf_ze_loopback PRIVATE ${ZE_LOADER_INCLUDE_DIRS} ${UMF_UTILS_DIR}
-                            ${UMF_SRC_DIR} ${UMF_CMAKE_SOURCE_DIR}/include)
+    add_library(umf_ze_loopback SHARED ze_loopback.h ze_loopback.cpp)
+    target_include_directories(
+        umf_ze_loopback PRIVATE ${ZE_LOADER_INCLUDE_DIRS} ${UMF_UTILS_DIR}
+                                ${UMF_SRC_DIR} ${UMF_CMAKE_SOURCE_DIR}/include)
+endif()

test/common/level_zero_mocks.cpp

@@ -1,3 +1,12 @@
+/*
+ *
+ * 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
+ *
+ */
+
 #include "level_zero_mocks.h"
 #include "umf/providers/provider_level_zero.h"
 #include <cstdlib>

test/common/level_zero_mocks.h

@@ -1,3 +1,12 @@
+/*
+ *
+ * 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
+ *
+ */
+
 #ifndef UMF_PROVIDER_LEVEL_ZERO_MOCKS_H
 #define UMF_PROVIDER_LEVEL_ZERO_MOCKS_H
 

test/common/ze_loopback.cpp

@@ -1,3 +1,12 @@
+/*
+ *
+ * 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
+ *
+ */
+
 #include "ze_loopback.h"
 #include "utils_log.h"
 

test/pools/pool_residency.cpp

@@ -126,6 +126,25 @@ TEST_F(PoolResidencyTestFixture,
     umfPoolFree(pool, ptr);
 }
 
+TEST_F(PoolResidencyTestFixture,
+       allocationThatFailedToBeMadeResidedShouldBeFreed) {
+    initializeMemoryPool(l0mock.initializeMemoryProviderWithResidentDevices(
+        OUR_DEVICE, {DEVICE_2}));
+
+    EXPECT_CALL(l0mock, zeMemAllocDevice(CONTEXT, _, _, _, OUR_DEVICE, _))
+        .WillOnce(
+            DoAll(SetArgPointee<5>(POINTER_0), Return(ZE_RESULT_SUCCESS)));
+    EXPECT_CALL(l0mock, zeContextMakeMemoryResident(CONTEXT, DEVICE_2, _, _))
+        .WillOnce(Return(ZE_RESULT_ERROR_DEVICE_LOST));
+    EXPECT_CALL(l0mock, zeMemFree(CONTEXT, _))
+        .WillOnce(Return(ZE_RESULT_ERROR_DEVICE_IN_LOW_POWER_STATE));
+
+    void *ptr = umfPoolMalloc(pool, 16 * 1024 * 1024);
+    EXPECT_EQ(ptr, nullptr);
+
+    umfPoolFree(pool, ptr);
+}
+
 int main(int argc, char **argv) {
     InitGoogleTest(&argc, argv);
     AddGlobalTestEnvironment(new MockedLevelZeroTestEnvironment);