From 32bdb08c860d5080b234fbae8da596165ef84218 Mon Sep 17 00:00:00 2001 From: Lukasz Dorau Date: Wed, 4 Nov 2026 08:23:54 +0000 Subject: [PATCH] [UR][L0v2] Bound retained kernels with capture-safe synchronization Immediate queues retained all submitted kernels until queue finish, keeping their programs alive as distinct kernels accumulated. Limit ordinary retention to 10,000 distinct kernels per immediate command list. Before the next launch, synchronize the list and release its pending kernels. Keep the command-list lock, but synchronize before kernel locking to avoid cross-queue deadlocks. In the fallback path, hold the kernel mutex continuously through argument setup and submission to prevent another queue from overwriting the launch arguments. Keep captured kernels separate and classify capture after submission to handle implicit fork/join capture. Skip limit cleanup during capture and preserve deferred retention for replayable command lists. Propagate capture-query errors before submission; afterward, log errors and retain kernels conservatively without reporting a submitted launch as failed. Add regression coverage for the limit, kernel reuse, explicit capture with full pending retention, implicit capture, and graph replay after cleanup. Exercise both launch paths with cross-queue host-gate waits and distinct arguments to detect deadlocks and argument corruption. Fixes: Jira URT-1346 Signed-off-by: Lukasz Dorau --- .../level_zero/v2/command_list_manager.cpp | 118 ++++-- .../level_zero/v2/command_list_manager.hpp | 9 +- .../v2/queue_immediate_in_order.cpp | 5 +- .../v2/queue_immediate_out_of_order.cpp | 4 +- .../level_zero/v2/deferred_kernel.cpp | 362 ++++++++++++++++++ 5 files changed, 466 insertions(+), 32 deletions(-) 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;