diff --git a/unified-runtime/source/adapters/level_zero/v2/command_list_manager.cpp b/unified-runtime/source/adapters/level_zero/v2/command_list_manager.cpp index 8c97ade034533..f4b7e97f9700e 100644 --- a/unified-runtime/source/adapters/level_zero/v2/command_list_manager.cpp +++ b/unified-runtime/source/adapters/level_zero/v2/command_list_manager.cpp @@ -114,8 +114,8 @@ void wait_list_view::addEvent(ur_event_handle_t phEvent) { ur_command_list_manager::ur_command_list_manager( ur_context_handle_t hContext, ur_device_handle_t hDevice, - v2::raii::command_list_unique_handle &&commandList) - : hContext(hContext), hDevice(hDevice), + v2::raii::command_list_unique_handle &&commandList, bool isImmediate) + : hContext(hContext), hDevice(hDevice), isImmediate(isImmediate), zeCommandList(std::move(commandList)) {} v2::raii::command_list_unique_handle && @@ -231,6 +231,29 @@ ur_command_list_manager::getSignalEvent(ur_event_handle_t hUserEvent, } } +// The command-list lock must be held, but no kernel mutex may be held while +// waiting for work that another queue may need to unblock. +ur_result_t ur_command_list_manager::releasePendingKernelsOnLimit() { + constexpr size_t maxPendingKernels = 10000; + if (isImmediate && pendingKernels.size() >= maxPendingKernels) { + bool captureActive = false; + auto result = queryGraphCaptureActive(&captureActive); + if (result != UR_RESULT_SUCCESS && + result != UR_RESULT_ERROR_UNSUPPORTED_FEATURE) + return result; + if (!captureActive) { + // Synchronize before launching so a failure cannot cause a submitted + // kernel to be retried. Captured kernels must remain available for + // replay. + ZE2UR_CALL(zeCommandListHostSynchronize, + (getZeCommandList(), UINT64_MAX)); + UR_CALL(releasePendingKernels()); + } + } + + return UR_RESULT_SUCCESS; +} + // must be called with hKernel->Mutex held ur_result_t ur_command_list_manager::appendKernelLaunchLocked( ur_kernel_handle_t hKernel, ze_kernel_handle_t hZeKernel, uint32_t workDim, @@ -276,7 +299,19 @@ ur_result_t ur_command_list_manager::appendKernelLaunchLocked( zeSignalEvent, waitListView.num, waitListView.handles)); } - recordSubmittedKernel(hKernel); + bool captureActive = false; + if (isImmediate) { + // A wait-list dependency can make this launch join another list's capture. + // Query after submission, but never report a submitted launch as failed. + auto result = queryGraphCaptureActive(&captureActive); + if (result != UR_RESULT_SUCCESS && + result != UR_RESULT_ERROR_UNSUPPORTED_FEATURE) { + UR_LOG(ERR, "Failed to query graph capture after kernel launch: {}", + result); + captureActive = true; + } + } + recordSubmittedKernel(hKernel, captureActive); postSubmit(hZeKernel, pGlobalWorkOffset); @@ -304,6 +339,8 @@ ur_result_t ur_command_list_manager::appendKernelLaunchUnlocked( return checkResult; } + UR_CALL(releasePendingKernelsOnLimit()); + ze_kernel_handle_t hZeKernel = hKernel->getZeHandle(hDevice.get()); std::scoped_lock Lock(hKernel->Mutex); @@ -313,21 +350,15 @@ ur_result_t ur_command_list_manager::appendKernelLaunchUnlocked( pLocalWorkSize, waitListView, phEvent, cooperative); } -ur_result_t ur_command_list_manager::appendKernelLaunch( - ur_kernel_handle_t hKernel, uint32_t workDim, - const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize, - const size_t *pLocalWorkSize, +static ur_result_t kernelLaunchPropertyChecks( const ur_kernel_launch_ext_properties_t *launchPropList, - wait_list_view &waitListView, ur_event_handle_t phEvent) { - TRACK_SCOPE_LATENCY("ur_command_list_manager::appendKernelLaunch"); - + bool &cooperative) { + cooperative = false; ur_kernel_launch_ext_properties_t *_launchPropList = const_cast(launchPropList); if (_launchPropList && _launchPropList->flags & UR_KERNEL_LAUNCH_FLAG_COOPERATIVE) { - UR_CALL(appendKernelLaunchUnlocked( - hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize, - waitListView, phEvent, true /* cooperative */)); + cooperative = true; return UR_RESULT_SUCCESS; } @@ -347,13 +378,24 @@ ur_result_t ur_command_list_manager::appendKernelLaunch( _launchPropList->pNext); } - UR_CALL(appendKernelLaunchUnlocked( - hKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize, - waitListView, phEvent, false /* cooperative */)); - return UR_RESULT_SUCCESS; } +ur_result_t ur_command_list_manager::appendKernelLaunch( + ur_kernel_handle_t hKernel, uint32_t workDim, + const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize, + const size_t *pLocalWorkSize, + const ur_kernel_launch_ext_properties_t *launchPropList, + wait_list_view &waitListView, ur_event_handle_t phEvent) { + TRACK_SCOPE_LATENCY("ur_command_list_manager::appendKernelLaunch"); + + bool cooperative = false; + UR_CALL(kernelLaunchPropertyChecks(launchPropList, cooperative)); + return appendKernelLaunchUnlocked(hKernel, workDim, pGlobalWorkOffset, + pGlobalWorkSize, pLocalWorkSize, + waitListView, phEvent, cooperative); +} + ur_result_t ur_command_list_manager::appendUSMMemcpy( bool blocking, void *pDst, const void *pSrc, size_t size, wait_list_view &waitListView, ur_event_handle_t phEvent) { @@ -1072,9 +1114,22 @@ ur_result_t ur_command_list_manager::appendNativeCommandExp( return UR_RESULT_ERROR_UNSUPPORTED_FEATURE; } -void ur_command_list_manager::recordSubmittedKernel( - ur_kernel_handle_t hKernel) { - auto [_, inserted] = submittedKernels.insert(hKernel); +ur_result_t ur_command_list_manager::releasePendingKernels() { + while (!pendingKernels.empty()) { + auto kernel = *pendingKernels.begin(); + pendingKernels.erase(pendingKernels.begin()); + UR_CALL(kernel->release()); + } + return UR_RESULT_SUCCESS; +} + +void ur_command_list_manager::recordSubmittedKernel(ur_kernel_handle_t hKernel, + bool captureActive) { + // Regular lists can be replayed, and captured kernels belong to a graph, + // not to the immediately executing command stream. + const bool canRetire = isImmediate && !captureActive; + auto &kernels = canRetire ? pendingKernels : submittedKernels; + auto [_, inserted] = kernels.insert(hKernel); if (inserted) { hKernel->RefCount.retain(); } @@ -1120,6 +1175,8 @@ ur_result_t ur_command_list_manager::appendEventsWaitWithBarrier( } ur_result_t ur_command_list_manager::releaseSubmittedKernels() { + UR_CALL(releasePendingKernels()); + // Free deferred kernels for (auto &hKernel : submittedKernels) { UR_CALL(hKernel->release()); @@ -1135,6 +1192,12 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExpOld( const ur_exp_kernel_arg_properties_t *pArgs, const ur_kernel_launch_ext_properties_t *launchPropList, wait_list_view &waitListView, ur_event_handle_t phEvent) { + UR_CALL(kernelLaunchChecks(hKernel, workDim)); + bool cooperative = false; + UR_CALL(kernelLaunchPropertyChecks(launchPropList, cooperative)); + UR_CALL(releasePendingKernelsOnLimit()); + ze_kernel_handle_t hZeKernel = hKernel->getZeHandle(hDevice.get()); + { std::scoped_lock guard(hKernel->Mutex); ur_device_handle_t hDevice = this->hDevice.get(); @@ -1171,13 +1234,11 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExpOld( return UR_RESULT_ERROR_INVALID_ENUMERATION; } } - } - UR_CALL(appendKernelLaunch(hKernel, workDim, pGlobalWorkOffset, - pGlobalWorkSize, pLocalWorkSize, launchPropList, - waitListView, phEvent)); - - return UR_RESULT_SUCCESS; + return appendKernelLaunchLocked( + hKernel, hZeKernel, workDim, pGlobalWorkOffset, pGlobalWorkSize, + pLocalWorkSize, waitListView, phEvent, cooperative); + } } ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExpNew( @@ -1232,6 +1293,8 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExpNew( _launchPropList->pNext); } + UR_CALL(releasePendingKernelsOnLimit()); + ze_kernel_handle_t hZeKernel = hKernel->getZeHandle(hDevice.get()); std::scoped_lock Lock(hKernel->Mutex); @@ -1434,6 +1497,9 @@ ur_result_t ur_command_list_manager::queryGraphCaptureActive(bool *pResult) { ze_result_t ZeResult = ZeGraphExt.normalizeGraphQueryResult(ZE_CALL_NOCHECK( ZeGraphExt.zeCommandListIsGraphCaptureEnabledExp, (getZeCommandList()))); + if (ZeResult != ZE_RESULT_QUERY_TRUE && ZeResult != ZE_RESULT_QUERY_FALSE) + return ze2urResult(ZeResult); + *pResult = (ZeResult == ZE_RESULT_QUERY_TRUE); return UR_RESULT_SUCCESS; diff --git a/unified-runtime/source/adapters/level_zero/v2/command_list_manager.hpp b/unified-runtime/source/adapters/level_zero/v2/command_list_manager.hpp index 35cd8e16a39ea..1337f5da51a9c 100644 --- a/unified-runtime/source/adapters/level_zero/v2/command_list_manager.hpp +++ b/unified-runtime/source/adapters/level_zero/v2/command_list_manager.hpp @@ -49,7 +49,8 @@ struct wait_list_view { struct ur_command_list_manager { ur_command_list_manager(ur_context_handle_t hContext, ur_device_handle_t hDevice, - v2::raii::command_list_unique_handle &&commandList); + v2::raii::command_list_unique_handle &&commandList, + bool isImmediate = false); ur_command_list_manager(const ur_command_list_manager &src) = delete; ur_command_list_manager(ur_command_list_manager &&src) = default; @@ -269,7 +270,9 @@ struct ur_command_list_manager { ur_event_handle_t phEvent, wait_list_view &waitListView, ur_command_t callerCommand); - void recordSubmittedKernel(ur_kernel_handle_t hKernel); + ur_result_t releasePendingKernelsOnLimit(); + ur_result_t releasePendingKernels(); + void recordSubmittedKernel(ur_kernel_handle_t hKernel, bool captureActive); ze_event_handle_t getSignalEvent(ur_event_handle_t hUserEvent, ur_command_t commandType); @@ -315,6 +318,8 @@ struct ur_command_list_manager { v2::raii::ur_device_handle_t hDevice; std::unordered_set submittedKernels; + bool isImmediate; + std::unordered_set pendingKernels; v2::raii::command_list_unique_handle zeCommandList; std::vector waitList; }; diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.cpp b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.cpp index 9e55b4edd160e..91025fa3623b4 100644 --- a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.cpp +++ b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.cpp @@ -34,7 +34,8 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t( hDevice->ZeDevice, {true, ordinal, true /* always enable copy offload */}, ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, priority, index, - priorityLevel)), + priorityLevel), + true), flags(flags), eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate) .borrow(hDevice->Id.value(), eventFlags)) {} @@ -44,7 +45,7 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t( raii::command_list_unique_handle commandListHandle, event_flags_t eventFlags, ur_queue_flags_t flags) : hContext(hContext), hDevice(hDevice), - commandListManager(hContext, hDevice, std::move(commandListHandle)), + commandListManager(hContext, hDevice, std::move(commandListHandle), true), flags(flags), eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate) .borrow(hDevice->Id.value(), eventFlags)) {} diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.cpp b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.cpp index 5d991fa608ac3..10e7dce4a7ac1 100644 --- a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.cpp +++ b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.cpp @@ -25,8 +25,8 @@ std::array createCommandListManagers( hContext->getCommandListCache().getImmediateCommandList( hDevice->ZeDevice, {true, ordinal, true /* always enable copy offload */}, - ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, priority, index, - priorityLevel)); + ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, priority, index, priorityLevel), + true); }); } diff --git a/unified-runtime/test/adapters/level_zero/v2/deferred_kernel.cpp b/unified-runtime/test/adapters/level_zero/v2/deferred_kernel.cpp index 675ed655e850b..893b378b521b4 100644 --- a/unified-runtime/test/adapters/level_zero/v2/deferred_kernel.cpp +++ b/unified-runtime/test/adapters/level_zero/v2/deferred_kernel.cpp @@ -4,9 +4,13 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // RUN: %with-v2 ./deferred_kernel-test +// RUN: %with-v2 env UR_L0_V2_DISABLE_ZE_LAUNCH_KERNEL_WITH_ARGS=1 ./deferred_kernel-test --gtest_filter=*SynchronizeAtKernelLimit* // REQUIRES: v2 // UNSUPPORTED: system-windows +#include +#include +#include #include #include "../../../conformance/enqueue/helpers.h" @@ -68,6 +72,364 @@ TEST_P(urEnqueueKernelLaunchTest, DeferredKernelRelease) { ASSERT_SUCCESS(urEventRelease(event)); } +struct urKernelReclamationTest : uur::urKernelExecutionTest { + void SetUp() override { + ASSERT_EQ(zeInit(ZE_INIT_FLAG_GPU_ONLY), ZE_RESULT_SUCCESS); + program_name = "foo"; + UUR_RETURN_ON_FATAL_FAILURE(urKernelExecutionTest::SetUp()); + } +}; +INSTANTIATE_TEST_SUITE_P( + , urKernelReclamationTest, + testing::Combine( + testing::ValuesIn(uur::DevicesEnvironment::instance->devices), + testing::Values(UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE, + static_cast( + UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE | + UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE))), + uur::devicePrinter); + +TEST_P(urKernelReclamationTest, SynchronizeAtKernelLimit) { + size_t global_size = 1; + size_t global_offset = 0; + const size_t numCommandLists = + std::get<1>(GetParam()) & UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE ? 4 + : 1; + auto zeGate = createZeEvent(context, device); + uur::raii::Event gate; + ASSERT_SUCCESS(urEventCreateWithNativeHandle( + reinterpret_cast(zeGate.get()), context, nullptr, + gate.ptr())); + ur_queue_properties_t properties = {UR_STRUCTURE_TYPE_QUEUE_PROPERTIES, + nullptr, + UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE}; + uur::raii::Queue otherQueue; + ASSERT_SUCCESS(urQueueCreate(context, device, &properties, otherQueue.ptr())); + + // Fill each immediate list with 10,000 distinct kernels without blocking GPU + // progress on a host signal. + std::vector kernels(10000); + for (auto &submittedKernel : kernels) { + if (&submittedKernel == &kernels.back()) { + // Reusing a kernel must not count towards the distinct-kernel limit. + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, kernels.front(), 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + uint32_t references = 0; + ASSERT_SUCCESS(urKernelGetInfo(kernels[1], UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, numCommandLists + 1); + } + ASSERT_SUCCESS( + urKernelCreate(program, kernel_name.c_str(), submittedKernel.ptr())); + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, submittedKernel, 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + } + + uint32_t references = 0; + ASSERT_SUCCESS(urKernelGetInfo(kernels.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, numCommandLists + 1); + + // The next launch on each list synchronizes and clears its retained kernels, + // including the previous reference to the kernel being reused. + ASSERT_SUCCESS(urEnqueueEventsWaitWithBarrier(queue, 1, gate.ptr(), nullptr)); + auto launch = std::async(std::launch::async, [&] { + return urEnqueueKernelLaunchWithArgsExp( + queue, kernels.front(), 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr); + }); + EXPECT_EQ(launch.wait_for(std::chrono::milliseconds(50)), + std::future_status::timeout); + // Another queue must be able to enqueue the same kernel before the host + // signals the gate. Signal even on timeout so a regression cannot hang. + auto otherLaunch = std::async(std::launch::async, [&] { + return urEnqueueKernelLaunchWithArgsExp( + otherQueue, kernels.front(), 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr); + }); + EXPECT_EQ(otherLaunch.wait_for(std::chrono::seconds(1)), + std::future_status::ready); + EXPECT_EQ(zeEventHostSignal(zeGate.get()), ZE_RESULT_SUCCESS); + EXPECT_SUCCESS(otherLaunch.get()); + ASSERT_SUCCESS(launch.get()); + ASSERT_SUCCESS(urQueueFinish(otherQueue)); + for (size_t i = 1; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, kernels.front(), 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + EXPECT_SUCCESS(urKernelGetInfo(kernels.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, numCommandLists + 1); + EXPECT_SUCCESS(urKernelGetInfo(kernels[1], UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 1u); + EXPECT_SUCCESS(urKernelGetInfo(kernels.back(), UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 1u); + + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, kernels.front(), 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + EXPECT_SUCCESS(urKernelGetInfo(kernels.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, numCommandLists + 1); + ASSERT_SUCCESS(urQueueFinish(queue)); + EXPECT_SUCCESS(urKernelGetInfo(kernels.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 1u); +} + +TEST_P(urKernelReclamationTest, SynchronizeAtKernelLimitPreservesArguments) { + constexpr size_t maxPendingKernels = 10000; + const size_t numCommandLists = + std::get<1>(GetParam()) & UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE ? 4 + : 1; + size_t global_size = 1; + size_t global_offset = 0; + + std::shared_ptr> fillBinary; + UUR_RETURN_ON_FATAL_FAILURE(uur::KernelsEnvironment::instance->LoadSource( + "fill", platform, fillBinary)); + uur::raii::Program fillProgram; + UUR_RETURN_ON_FATAL_FAILURE(uur::KernelsEnvironment::instance->CreateProgram( + platform, context, device, *fillBinary, nullptr, fillProgram.ptr())); + ASSERT_SUCCESS(urProgramBuild(context, fillProgram, nullptr)); + const auto fillNames = + uur::KernelsEnvironment::instance->GetEntryPointNames("fill"); + ASSERT_FALSE(fillNames.empty()); + uur::raii::Kernel sharedKernel; + ASSERT_SUCCESS(urKernelCreate(fillProgram, fillNames.front().c_str(), + sharedKernel.ptr())); + + ur_queue_properties_t properties = {UR_STRUCTURE_TYPE_QUEUE_PROPERTIES, + nullptr, + UR_QUEUE_FLAG_SUBMISSION_IMMEDIATE}; + uur::raii::Queue otherQueue; + ASSERT_SUCCESS(urQueueCreate(context, device, &properties, otherQueue.ptr())); + ur_queue_handle_t otherQueueHandle = otherQueue; + ur_kernel_handle_t sharedKernelHandle = sharedKernel; + uur::KernelLaunchHelper firstArgs(platform, context, sharedKernelHandle, + queue); + uur::KernelLaunchHelper secondArgs(platform, context, sharedKernelHandle, + otherQueueHandle); + uur::raii::Mem firstBuffer; + uur::raii::Mem secondBuffer; + constexpr uint32_t firstValue = 42; + constexpr uint32_t secondValue = 43; + UUR_RETURN_ON_FATAL_FAILURE( + firstArgs.AddBuffer1DArg(sizeof(firstValue), firstBuffer.ptr(), nullptr)); + firstArgs.AddPodArg(firstValue); + UUR_RETURN_ON_FATAL_FAILURE(secondArgs.AddBuffer1DArg( + sizeof(secondValue), secondBuffer.ptr(), nullptr)); + secondArgs.AddPodArg(secondValue); + + std::vector kernels(maxPendingKernels); + for (auto &submittedKernel : kernels) { + ASSERT_SUCCESS( + urKernelCreate(program, kernel_name.c_str(), submittedKernel.ptr())); + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, submittedKernel, 1, &global_offset, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + } + + auto zeGate = createZeEvent(context, device); + uur::raii::Event gate; + ASSERT_SUCCESS(urEventCreateWithNativeHandle( + reinterpret_cast(zeGate.get()), context, nullptr, + gate.ptr())); + ASSERT_SUCCESS(urEnqueueEventsWaitWithBarrier(queue, 1, gate.ptr(), nullptr)); + auto firstLaunch = std::async(std::launch::async, [&] { + return urEnqueueKernelLaunchWithArgsExp( + queue, sharedKernel, 1, &global_offset, &global_size, nullptr, + firstArgs.GetNumArgs(), firstArgs.GetArgs(), nullptr, 0, nullptr, + nullptr); + }); + EXPECT_EQ(firstLaunch.wait_for(std::chrono::milliseconds(50)), + std::future_status::timeout); + auto secondLaunch = std::async(std::launch::async, [&] { + return urEnqueueKernelLaunchWithArgsExp( + otherQueue, sharedKernel, 1, &global_offset, &global_size, nullptr, + secondArgs.GetNumArgs(), secondArgs.GetArgs(), nullptr, 0, nullptr, + nullptr); + }); + // Always release the gate so a locking regression fails without hanging. + EXPECT_EQ(secondLaunch.wait_for(std::chrono::seconds(1)), + std::future_status::ready); + EXPECT_EQ(zeEventHostSignal(zeGate.get()), ZE_RESULT_SUCCESS); + EXPECT_SUCCESS(secondLaunch.get()); + EXPECT_SUCCESS(firstLaunch.get()); + ASSERT_SUCCESS(urQueueFinish(queue)); + ASSERT_SUCCESS(urQueueFinish(otherQueue)); + uint32_t firstResult = 0; + uint32_t secondResult = 0; + ASSERT_SUCCESS(urEnqueueMemBufferRead(queue, firstBuffer, true, 0, + sizeof(firstResult), &firstResult, 0, + nullptr, nullptr)); + ASSERT_SUCCESS(urEnqueueMemBufferRead(otherQueue, secondBuffer, true, 0, + sizeof(secondResult), &secondResult, 0, + nullptr, nullptr)); + EXPECT_EQ(firstResult, firstValue); + EXPECT_EQ(secondResult, secondValue); +} + +TEST_P(urKernelReclamationTest, CapturedKernelsAreNotSynchronizedAtLimit) { + constexpr size_t maxPendingKernels = 10000; + ur_bool_t supported = false; + ASSERT_SUCCESS(urDeviceGetInfo( + device, UR_DEVICE_INFO_GRAPH_RECORD_AND_REPLAY_SUPPORT_EXP, + sizeof(supported), &supported, nullptr)); + if (!supported) + GTEST_SKIP() << "Graph capture is not supported."; + + const size_t numCommandLists = + (getQueueMode() & UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE) ? 4 : 1; + size_t global_size = 1; + std::vector pending(maxPendingKernels); + for (auto &submittedKernel : pending) { + ASSERT_SUCCESS( + urKernelCreate(program, kernel_name.c_str(), submittedKernel.ptr())); + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, submittedKernel, 1, nullptr, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + } + + std::vector kernels(maxPendingKernels + 1); + for (auto &capturedKernel : kernels) { + ASSERT_SUCCESS( + urKernelCreate(program, kernel_name.c_str(), capturedKernel.ptr())); + } + + ASSERT_SUCCESS(urQueueBeginGraphCaptureExp(queue)); + for (size_t repeat = 0; repeat < 2; ++repeat) { + for (auto &capturedKernel : kernels) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, capturedKernel, 1, nullptr, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + } + uint32_t references = 0; + ASSERT_SUCCESS(urKernelGetInfo(pending.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, numCommandLists + 1); + ur_exp_graph_handle_t graph = nullptr; + ASSERT_SUCCESS(urQueueEndGraphCaptureExp(queue, &graph)); + + // Ordinary submissions can retire pending kernels, but not captured ones. + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, kernel, 1, nullptr, &global_size, nullptr, GetNumArgs(), + GetArgs(), nullptr, 0, nullptr, nullptr)); + } + ASSERT_SUCCESS(urKernelGetInfo(pending.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 1u); + ASSERT_SUCCESS(urKernelGetInfo(kernels.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 2u); + + ur_exp_executable_graph_handle_t executable = nullptr; + ASSERT_SUCCESS(urGraphInstantiateGraphExp(graph, &executable)); + ASSERT_SUCCESS(urEnqueueGraphExp(queue, executable, 0, nullptr, nullptr)); + uur::raii::Event completed; + ASSERT_SUCCESS( + urEnqueueEventsWaitWithBarrier(queue, 0, nullptr, completed.ptr())); + ASSERT_SUCCESS(urEventWait(1, completed.ptr())); + ASSERT_SUCCESS(urKernelGetInfo(kernels.front(), + UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 2u); + ASSERT_SUCCESS(urGraphExecutableGraphDestroyExp(executable)); + ASSERT_SUCCESS(urGraphDestroyExp(graph)); + ASSERT_SUCCESS(urQueueFinish(queue)); +} + +TEST_P(urKernelReclamationTest, ImplicitlyCapturedKernelSurvivesLimitCleanup) { + ur_bool_t supported = false; + ASSERT_SUCCESS(urDeviceGetInfo( + device, UR_DEVICE_INFO_GRAPH_RECORD_AND_REPLAY_SUPPORT_EXP, + sizeof(supported), &supported, nullptr)); + if (!supported) + GTEST_SKIP() << "Graph capture is not supported."; + + ur_queue_properties_t properties = {UR_STRUCTURE_TYPE_QUEUE_PROPERTIES, + nullptr, getQueueMode()}; + uur::raii::Queue primaryQueue; + ASSERT_SUCCESS( + urQueueCreate(context, device, &properties, primaryQueue.ptr())); + ASSERT_SUCCESS(urQueueBeginGraphCaptureExp(primaryQueue)); + uur::raii::Event fork; + ASSERT_SUCCESS( + urEnqueueEventsWaitWithBarrier(primaryQueue, 0, nullptr, fork.ptr())); + + size_t global_size = 1; + uur::raii::Event joined; + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, kernel, 1, nullptr, &global_size, nullptr, GetNumArgs(), GetArgs(), + nullptr, 1, fork.ptr(), joined.ptr())); + bool capturing = false; + ASSERT_SUCCESS(urQueueIsGraphCaptureEnabledExp(queue, &capturing)); + ASSERT_TRUE(capturing); + ASSERT_SUCCESS( + urEnqueueEventsWaitWithBarrier(primaryQueue, 1, joined.ptr(), nullptr)); + ur_exp_graph_handle_t graph = nullptr; + ASSERT_SUCCESS(urQueueEndGraphCaptureExp(primaryQueue, &graph)); + + const size_t numCommandLists = + (getQueueMode() & UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE) ? 4 : 1; + std::vector kernels(10000); + for (auto &submittedKernel : kernels) { + ASSERT_SUCCESS( + urKernelCreate(program, kernel_name.c_str(), submittedKernel.ptr())); + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, submittedKernel, 1, nullptr, &global_size, nullptr, + GetNumArgs(), GetArgs(), nullptr, 0, nullptr, nullptr)); + } + } + for (size_t i = 0; i < numCommandLists; ++i) { + ASSERT_SUCCESS(urEnqueueKernelLaunchWithArgsExp( + queue, kernels.front(), 1, nullptr, &global_size, nullptr, GetNumArgs(), + GetArgs(), nullptr, 0, nullptr, nullptr)); + } + + uint32_t references = 0; + ASSERT_SUCCESS(urKernelGetInfo(kernel, UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 2u); + ASSERT_SUCCESS(urKernelGetInfo(kernels.back(), UR_KERNEL_INFO_REFERENCE_COUNT, + sizeof(references), &references, nullptr)); + EXPECT_EQ(references, 1u); + + ur_exp_executable_graph_handle_t executable = nullptr; + ASSERT_SUCCESS(urGraphInstantiateGraphExp(graph, &executable)); + ASSERT_SUCCESS( + urEnqueueGraphExp(primaryQueue, executable, 0, nullptr, nullptr)); + ASSERT_SUCCESS(urQueueFinish(primaryQueue)); + ASSERT_SUCCESS(urGraphExecutableGraphDestroyExp(executable)); + ASSERT_SUCCESS(urGraphDestroyExp(graph)); + ASSERT_SUCCESS(urQueueFinish(queue)); +} + struct urMultiQueueLaunchKernelDeferFreeTest : uur::urMultiQueueMultiDeviceTest<2> { std::string KernelName;