Explicitly specify USMFree as a blocking operation.

Also add test to enforce this.

Author: aarongreig

include/ur_api.h

@@ -3661,6 +3661,11 @@ urUSMSharedAlloc(
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Free the USM memory object
 ///
+/// @details
+///     - Note that implementations are required to wait for previously enqueued
+///       commands that may be accessing `pMem` to finish before freeing the
+///       memory.
+///
 /// @returns
 ///     - ::UR_RESULT_SUCCESS
 ///     - ::UR_RESULT_ERROR_UNINITIALIZED

scripts/core/usm.yml

@@ -364,6 +364,8 @@ desc: "Free the USM memory object"
 class: $xUSM
 name: Free
 ordinal: "0"
+details:
+  - "Note that implementations are required to wait for previously enqueued commands that may be accessing `pMem` to finish before freeing the memory."
 params:
     - type: $x_context_handle_t
       name: hContext

source/loader/ur_libapi.cpp

@@ -2404,6 +2404,11 @@ ur_result_t UR_APICALL urUSMSharedAlloc(
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Free the USM memory object
 ///
+/// @details
+///     - Note that implementations are required to wait for previously enqueued
+///       commands that may be accessing `pMem` to finish before freeing the
+///       memory.
+///
 /// @returns
 ///     - ::UR_RESULT_SUCCESS
 ///     - ::UR_RESULT_ERROR_UNINITIALIZED

source/ur_api.cpp

@@ -2080,6 +2080,11 @@ ur_result_t UR_APICALL urUSMSharedAlloc(
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Free the USM memory object
 ///
+/// @details
+///     - Note that implementations are required to wait for previously enqueued
+///       commands that may be accessing `pMem` to finish before freeing the
+///       memory.
+///
 /// @returns
 ///     - ::UR_RESULT_SUCCESS
 ///     - ::UR_RESULT_ERROR_UNINITIALIZED

test/conformance/CMakeLists.txt

@@ -112,7 +112,6 @@ add_subdirectory(platform)
 add_subdirectory(device)
 add_subdirectory(context)
 add_subdirectory(memory)
-add_subdirectory(usm)
 add_subdirectory(event)
 add_subdirectory(queue)
 add_subdirectory(sampler)
@@ -129,6 +128,7 @@ set(TEST_SUBDIRECTORIES_DPCXX
     "exp_usm_p2p"
     "exp_launch_properties"
     "memory-migrate"
+    "usm"
 )
 
 if(UR_DPCXX)

test/conformance/usm/CMakeLists.txt

@@ -3,7 +3,7 @@
 # See LICENSE.TXT
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-add_conformance_test_with_devices_environment(usm 
+add_conformance_test_with_kernels_environment(usm
     urUSMDeviceAlloc.cpp
     urUSMFree.cpp
     urUSMGetMemAllocInfo.cpp

test/conformance/usm/urUSMFree.cpp

@@ -95,3 +95,84 @@ TEST_P(urUSMFreeTest, InvalidNullPtrMem) {
     ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_POINTER,
                      urUSMFree(context, nullptr));
 }
+
+// This goal of this test is to ensure urUSMFree blocks and waits for operations
+// accessing the given allocation to finish before actually freeing the memory.
+struct urUSMFreeDuringExecutionTest : uur::urKernelExecutionTest {
+    void SetUp() {
+        program_name = "fill_usm";
+        UUR_RETURN_ON_FATAL_FAILURE(urKernelExecutionTest::SetUp());
+    }
+
+    void *allocation = nullptr;
+    size_t array_size = 256;
+    size_t allocation_size = array_size * sizeof(uint32_t);
+    uint32_t data = 42;
+    size_t wg_offset = 0;
+};
+UUR_INSTANTIATE_KERNEL_TEST_SUITE_P(urUSMFreeDuringExecutionTest);
+
+TEST_P(urUSMFreeDuringExecutionTest, SuccessHost) {
+    ur_device_usm_access_capability_flags_t host_usm_flags = 0;
+    ASSERT_SUCCESS(uur::GetDeviceUSMHostSupport(device, host_usm_flags));
+    if (!(host_usm_flags & UR_DEVICE_USM_ACCESS_CAPABILITY_FLAG_ACCESS)) {
+        GTEST_SKIP() << "Host USM is not supported.";
+    }
+
+    ASSERT_SUCCESS(urUSMHostAlloc(context, nullptr, nullptr, allocation_size,
+                                  &allocation));
+    ASSERT_NE(allocation, nullptr);
+
+    EXPECT_SUCCESS(urKernelSetArgPointer(kernel, 0, nullptr, allocation));
+    EXPECT_SUCCESS(
+        urKernelSetArgValue(kernel, 1, sizeof(data), nullptr, &data));
+    EXPECT_SUCCESS(urEnqueueKernelLaunch(queue, kernel, 1, &wg_offset,
+                                         &array_size, nullptr, 0, nullptr,
+                                         nullptr));
+    ASSERT_SUCCESS(urUSMFree(context, allocation));
+    ASSERT_SUCCESS(urQueueFinish(queue));
+}
+
+TEST_P(urUSMFreeDuringExecutionTest, SuccessDevice) {
+    ur_device_usm_access_capability_flags_t device_usm_flags = 0;
+    ASSERT_SUCCESS(uur::GetDeviceUSMDeviceSupport(device, device_usm_flags));
+    if (!(device_usm_flags & UR_DEVICE_USM_ACCESS_CAPABILITY_FLAG_ACCESS)) {
+        GTEST_SKIP() << "Device USM is not supported.";
+    }
+
+    ASSERT_SUCCESS(urUSMDeviceAlloc(context, device, nullptr, nullptr,
+                                    allocation_size, &allocation));
+    ASSERT_NE(allocation, nullptr);
+
+    EXPECT_SUCCESS(urKernelSetArgPointer(kernel, 0, nullptr, allocation));
+    EXPECT_SUCCESS(
+        urKernelSetArgValue(kernel, 1, sizeof(data), nullptr, &data));
+
+    EXPECT_SUCCESS(urEnqueueKernelLaunch(queue, kernel, 1, &wg_offset,
+                                         &array_size, nullptr, 0, nullptr,
+                                         nullptr));
+    ASSERT_SUCCESS(urUSMFree(context, allocation));
+    ASSERT_SUCCESS(urQueueFinish(queue));
+}
+
+TEST_P(urUSMFreeDuringExecutionTest, SuccessShared) {
+    ur_device_usm_access_capability_flags_t shared_usm_flags = 0;
+    ASSERT_SUCCESS(
+        uur::GetDeviceUSMSingleSharedSupport(device, shared_usm_flags));
+    if (!(shared_usm_flags & UR_DEVICE_USM_ACCESS_CAPABILITY_FLAG_ACCESS)) {
+        GTEST_SKIP() << "Shared USM is not supported.";
+    }
+
+    ASSERT_SUCCESS(urUSMSharedAlloc(context, device, nullptr, nullptr,
+                                    allocation_size, &allocation));
+    ASSERT_NE(allocation, nullptr);
+
+    EXPECT_SUCCESS(urKernelSetArgPointer(kernel, 0, nullptr, allocation));
+    EXPECT_SUCCESS(
+        urKernelSetArgValue(kernel, 1, sizeof(data), nullptr, &data));
+    EXPECT_SUCCESS(urEnqueueKernelLaunch(queue, kernel, 1, &wg_offset,
+                                         &array_size, nullptr, 0, nullptr,
+                                         nullptr));
+    ASSERT_SUCCESS(urUSMFree(context, allocation));
+    ASSERT_SUCCESS(urQueueFinish(queue));
+}

test/conformance/usm/usm_adapter_native_cpu.match

@@ -3,6 +3,7 @@ urUSMDeviceAllocTest.InvalidUSMSize/*__UsePoolDisabled
 urUSMFreeTest.SuccessDeviceAlloc/*
 urUSMFreeTest.SuccessHostAlloc/*
 urUSMFreeTest.SuccessSharedAlloc/*
+urUSMFreeDuringExecutionTest.*
 urUSMGetMemAllocInfoTest.Success/*__UR_USM_ALLOC_INFO_TYPE
 urUSMGetMemAllocInfoTest.Success/*__UR_USM_ALLOC_INFO_BASE_PTR
 urUSMGetMemAllocInfoTest.Success/*__UR_USM_ALLOC_INFO_SIZE