From ed4aa1f25a4e2b1d81aa97a15b2fcab78e04447d Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 1 Oct 2026 15:34:05 -0700 Subject: [PATCH 01/46] Add D3D12 accelerated weight loading Load external WebGPU initializers through a reusable Win32 and D3D12 file loader and import the resulting resources into Dawn. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/workflows/windows_webgpu.yml | 3 +- cmake/CMakeLists.txt | 5 + cmake/onnxruntime_d3d12_file_loader.cmake | 28 + cmake/onnxruntime_providers_webgpu.cmake | 17 + js/common/lib/inference-session.ts | 12 + js/node/src/session_options_helper.cc | 1 + js/node/test/api/simple-api-tests.ts | 20 + .../core/framework/external_data_loader.h | 80 +- .../framework/external_data_loader_manager.cc | 54 +- .../framework/external_data_loader_manager.h | 8 + .../core/framework/session_state_utils.cc | 64 +- .../core/framework/tensorprotoutils.cc | 42 + onnxruntime/core/framework/tensorprotoutils.h | 5 + .../d3d12_file_buffer_loader.cc | 1067 +++++++++++++++++ .../d3d12_file_buffer_loader.h | 88 ++ .../core/providers/js/external_data_loader.cc | 2 + .../core/providers/js/external_data_loader.h | 2 + .../core/providers/webgpu/allocator.cc | 13 +- onnxruntime/core/providers/webgpu/allocator.h | 4 +- .../webgpu/d3d12_external_data_loader.cc | 916 ++++++++++++++ .../webgpu/d3d12_external_data_loader.h | 91 ++ .../core/providers/webgpu/data_transfer.cc | 25 +- .../core/providers/webgpu/data_transfer.h | 14 +- .../core/providers/webgpu/ep/sync_stream.cc | 2 +- .../providers/webgpu/external_data_loader.cc | 2 + .../providers/webgpu/external_data_loader.h | 2 + .../core/providers/webgpu/webgpu_context.cc | 690 +++++++---- .../core/providers/webgpu/webgpu_context.h | 41 +- .../webgpu/webgpu_execution_provider.cc | 57 +- .../webgpu/webgpu_execution_provider.h | 18 +- .../webgpu/webgpu_provider_factory.cc | 33 +- .../webgpu/webgpu_provider_options.h | 20 + onnxruntime/core/session/inference_session.cc | 1 + .../framework/external_data_loader_test.cc | 111 +- .../webgpu/concurrent_context_test.cc | 448 +++++++ .../providers/webgpu/webgpu_context_test.cc | 629 +++++++++- onnxruntime/test/webgpu/delay_load/main.cc | 3 +- onnxruntime/test/webgpu/external_dawn/main.cc | 3 + setup.py | 2 +- 39 files changed, 4317 insertions(+), 306 deletions(-) create mode 100644 cmake/onnxruntime_d3d12_file_loader.cmake create mode 100644 onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc create mode 100644 onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h create mode 100644 onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc create mode 100644 onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h diff --git a/.github/workflows/windows_webgpu.yml b/.github/workflows/windows_webgpu.yml index dd85d11ec4a1c..e2457ad51e2c0 100644 --- a/.github/workflows/windows_webgpu.yml +++ b/.github/workflows/windows_webgpu.yml @@ -134,7 +134,8 @@ jobs: ${{ matrix.vcpkg_option == 'vcpkg' && '--use_vcpkg --use_vcpkg_ms_internal_asset_cache' || '' }} ` --cmake_extra_defines ` onnxruntime_BUILD_UNIT_TESTS=ON ` - onnxruntime_BUILD_DAWN_SHARED_LIBRARY=ON + onnxruntime_BUILD_DAWN_SHARED_LIBRARY=ON ` + onnxruntime_ENABLE_D3D12_FILE_LOADING=ON if ($lastExitCode -ne 0) { exit $lastExitCode } diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 3a3bda1dded41..d7fa631df9750 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -230,6 +230,11 @@ if(onnxruntime_DAWN_PREBUILT_DIR) endif() option(onnxruntime_BUILD_DAWN_SHARED_LIBRARY "Build Dawn as a shared library" OFF) option(onnxruntime_ENABLE_PIX_FOR_WEBGPU_EP "Adding frame present for PIX to capture a frame" OFF) +option(onnxruntime_ENABLE_D3D12_FILE_LOADING + "Enable accelerated file-to-D3D12-buffer loading on Windows" OFF) +if(onnxruntime_ENABLE_D3D12_FILE_LOADING) + add_compile_definitions(ENABLE_D3D12_FILE_LOADING=1) +endif() # The following 2 options are only for Windows option(onnxruntime_ENABLE_DAWN_BACKEND_VULKAN "Enable Vulkan backend for Dawn (on Windows)" OFF) option(onnxruntime_ENABLE_DAWN_BACKEND_D3D12 "Enable D3D12 backend for Dawn (on Windows)" ON) diff --git a/cmake/onnxruntime_d3d12_file_loader.cmake b/cmake/onnxruntime_d3d12_file_loader.cmake new file mode 100644 index 0000000000000..8a55047679779 --- /dev/null +++ b/cmake/onnxruntime_d3d12_file_loader.cmake @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +include_guard(GLOBAL) + +if(NOT WIN32 OR NOT MSVC) + message(FATAL_ERROR "onnxruntime_d3d12_file_loader requires native Windows MSVC.") +endif() + +set(onnxruntime_d3d12_file_loader_srcs + "${ONNXRUNTIME_ROOT}/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h" + "${ONNXRUNTIME_ROOT}/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc" +) + +source_group(TREE ${ONNXRUNTIME_ROOT} FILES ${onnxruntime_d3d12_file_loader_srcs}) +onnxruntime_add_static_library( + onnxruntime_d3d12_file_loader + ${onnxruntime_d3d12_file_loader_srcs}) +onnxruntime_add_include_to_target( + onnxruntime_d3d12_file_loader + onnxruntime_common ${WIL_TARGET} ${GSL_TARGET}) +target_link_libraries( + onnxruntime_d3d12_file_loader + PUBLIC d3d12.lib dxgi.lib dxguid.lib + PRIVATE ${WIL_TARGET} ${GSL_TARGET}) +set_target_properties( + onnxruntime_d3d12_file_loader + PROPERTIES FOLDER "ONNXRuntime") diff --git a/cmake/onnxruntime_providers_webgpu.cmake b/cmake/onnxruntime_providers_webgpu.cmake index 40402e3b0638d..d5b2c6cb2029c 100644 --- a/cmake/onnxruntime_providers_webgpu.cmake +++ b/cmake/onnxruntime_providers_webgpu.cmake @@ -7,6 +7,17 @@ add_compile_definitions(USE_WEBGPU=1) + if(onnxruntime_ENABLE_D3D12_FILE_LOADING) + if(NOT WIN32 OR NOT MSVC OR NOT onnxruntime_ENABLE_DAWN_BACKEND_D3D12) + message(FATAL_ERROR "D3D12 accelerated weight loading requires native Windows MSVC and the Dawn D3D12 backend.") + endif() + if(onnxruntime_USE_EP_API_ADAPTERS) + message(FATAL_ERROR "D3D12 accelerated weight loading is not supported with EP API adapter/plugin builds.") + endif() + + include("${CMAKE_CURRENT_LIST_DIR}/onnxruntime_d3d12_file_loader.cmake") + endif() + if (onnxruntime_ENABLE_WEBASSEMBLY_THREADS) add_definitions(-DENABLE_WEBASSEMBLY_THREADS=1) endif() @@ -177,6 +188,12 @@ set_target_properties(onnxruntime_providers_webgpu PROPERTIES CXX_STANDARD_REQUIRED ON) set_target_properties(onnxruntime_providers_webgpu PROPERTIES FOLDER "ONNXRuntime") + if(onnxruntime_ENABLE_D3D12_FILE_LOADING) + target_link_libraries( + onnxruntime_providers_webgpu + PRIVATE onnxruntime_d3d12_file_loader) + endif() + if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten") # target "emdawnwebgpu_c" is created by Dawn, including "-fno-exceptions" in its compile options by default. # diff --git a/js/common/lib/inference-session.ts b/js/common/lib/inference-session.ts index 42d014e4f3c69..c7f665e784811 100644 --- a/js/common/lib/inference-session.ts +++ b/js/common/lib/inference-session.ts @@ -310,6 +310,18 @@ export declare namespace InferenceSession { */ enableRobustness?: boolean; + /** + * Select accelerated external-weight loading for the native Windows WebGPU execution provider. + * - 'off': Disable accelerated weight loading. + * - 'preferred': Use D3D12 accelerated loading when available and otherwise use the ordinary loading path. + * - 'required': Require D3D12 accelerated loading support. + * + * This setting is available only in ONNX Runtime (Node.js binding) builds with D3D12 file loading support. + * + * @default 'off' + */ + weightLoadAcceleration?: 'off' | 'preferred' | 'required'; + /** * Specify the cache mode for storage buffers. * - 'disabled': Disable buffer cache. Buffers are destroyed when no longer in use. diff --git a/js/node/src/session_options_helper.cc b/js/node/src/session_options_helper.cc index da55fb98af236..cbd008bf8701b 100644 --- a/js/node/src/session_options_helper.cc +++ b/js/node/src/session_options_helper.cc @@ -77,6 +77,7 @@ void ParseExecutionProviders(const Napi::Array epList, Ort::SessionOptions& sess std::string value; if (name == "preferredLayout" || name == "validationMode" || + name == "weightLoadAcceleration" || name == "storageBufferCacheMode" || name == "uniformBufferCacheMode" || name == "queryResolveBufferCacheMode" || diff --git a/js/node/test/api/simple-api-tests.ts b/js/node/test/api/simple-api-tests.ts index da7adecc3d4e4..0f34485d76d91 100644 --- a/js/node/test/api/simple-api-tests.ts +++ b/js/node/test/api/simple-api-tests.ts @@ -88,6 +88,26 @@ const MODEL_TEST_TYPES_CASES: Array<{ ]; describe('API Tests - simple API tests', () => { + it('recognizes every WebGPU weightLoadAcceleration value', async () => { + const values: ReadonlyArray> = + ['off', 'preferred', 'required']; + const model = path.join(TEST_DATA_ROOT, 'test_types_float.onnx'); + + for (const weightLoadAcceleration of values) { + try { + const session = await InferenceSession.create(model, { + executionProviders: [{ name: 'webgpu', weightLoadAcceleration }], + }); + await session.release(); + } catch (error) { + // Builds without WebGPU/D3D12 file loading and machines without compatible + // hardware can reject the provider after parsing. The regression here + // is specifically that Node must recognize and forward the option. + assert.doesNotMatch(String(error), /WebGPU EP has an unrecognized option: 'weightLoadAcceleration'/); + } + } + }); + MODEL_TEST_TYPES_CASES.forEach((testCase) => { it(`${testCase.model}`, async () => { const session = await InferenceSession.create(testCase.model); diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index 43549fd3e6610..535de3f9a29dc 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -4,16 +4,20 @@ #pragma once #include -#include #include +#include +#include +#include #include "core/common/common.h" #include "core/common/safeint.h" +#include "core/framework/ortdevice.h" #include "core/platform/env.h" struct OrtMemoryInfo; namespace onnxruntime { +class IAllocator; #ifndef SHARED_PROVIDER class Tensor; #endif @@ -24,12 +28,86 @@ class Status; } // Data transfer interface. +// Keep default virtual methods inline so shared providers do not require a non-exported RTTI symbol. class IExternalDataLoader { public: virtual ~IExternalDataLoader() = default; virtual bool CanLoad(const OrtMemoryInfo& target_memory_info) const = 0; + virtual bool SupportsDataType([[maybe_unused]] int32_t tensor_data_type) const { return true; } + +#if defined(ENABLE_D3D12_FILE_LOADING) + // Returns true when the loader creates the tensor's backing allocation instead of + // writing into a tensor allocated by the framework. + virtual bool CreatesTensorForDevice([[maybe_unused]] const OrtDevice& target_device) const { return false; } + + // Begins a batch covering the external initializers for one SessionState + // initialization. ORT describes all loader-created tensors through PrepareTensor() + // before calling FinalizeLoad(), allowing an implementation to coalesce file + // ranges, allocate destination resources together, and perform shared I/O or device + // work before individual tensors become visible through LoadTensor(). Loaders that + // can populate each tensor independently do not need to override the batch hooks. + // + // This is called once before zero or more PrepareTensor() calls and before any + // initializer is loaded through LoadTensor(). Implementations should discard stale + // batch state and initialize any batch-wide resources. Returning an error prevents + // the batch from continuing and causes AbortLoad() to be called. + virtual common::Status BeginLoad() const { return common::Status::OK(); } + + // Adds one external tensor to the current batch. The path, offset, and length + // have already been resolved and validated against the TensorProto by the caller. + // Implementations may record the range, reserve its destination, or start deferred + // work. ORT does not call LoadTensor() during this preparation phase. After + // FinalizeLoad() succeeds, ORT makes a separate initializer-deserialization pass + // and calls LoadTensor() for each tensor still claimed by this loader. Cancellation + // or another initialization error may prevent some or all of those calls. + // + // PrepareTensor() calls occur after a successful BeginLoad() and before + // FinalizeLoad(). The same file may be referenced by multiple tensors, and a batch + // may contain no tensors. Returning an error abandons the batch and causes + // AbortLoad() to be called. The default implementation performs no preparation. + virtual common::Status PrepareTensor([[maybe_unused]] const Env& env, + [[maybe_unused]] const std::filesystem::path& data_file_path, + [[maybe_unused]] std::string_view tensor_name, + [[maybe_unused]] FileOffsetType data_offset, + [[maybe_unused]] SafeInt data_length) const { + return common::Status::OK(); + } + + // Completes preparation for the current batch. On success, the loader must + // either make every tensor for which PrepareTensor() returned OK ready for a + // subsequent LoadTensor() call, or make CanLoad() return false for those tensors so + // normal loading can take over. Any asynchronous I/O or device work required to + // make tensors ready must be completed or synchronized here. The implementation + // must keep prepared tensors alive until LoadTensor() transfers their ownership or + // AbortLoad() releases them. This method is called even for an empty batch. + // Long-running work should check is_canceled and return MODEL_LOAD_CANCELED when + // cancellation is observed. Returning any error causes AbortLoad() to be called. + // The default implementation completes immediately. + virtual common::Status FinalizeLoad([[maybe_unused]] const std::function& is_canceled) const { + return common::Status::OK(); + } + + // Abandons the active batch. Cancel outstanding work and release resources + // still owned by the batch, without invalidating allocations already transferred by + // successful LoadTensor() calls. This may be called after BeginLoad(), + // PrepareTensor(), FinalizeLoad(), or partial tensor consumption, and may be called + // more than once. Implementations must therefore be noexcept and idempotent. + virtual void AbortLoad() const noexcept {} + + // A tensor-creating loader replaces tensor with one backed by memory owned through allocator. + virtual common::Status LoadTensor([[maybe_unused]] const Env& env, + [[maybe_unused]] const std::filesystem::path& data_file_path, + [[maybe_unused]] std::string_view tensor_name, + [[maybe_unused]] FileOffsetType data_offset, + [[maybe_unused]] SafeInt data_length, + [[maybe_unused]] const std::shared_ptr& allocator, + [[maybe_unused]] Tensor& tensor) const { + ORT_NOT_IMPLEMENTED(__FUNCTION__, " is not implemented"); + } +#endif + // Tensor should be already allocated with the correct memory info and size. #if defined(__wasm__) virtual common::Status LoadTensor([[maybe_unused]] const Env& env, diff --git a/onnxruntime/core/framework/external_data_loader_manager.cc b/onnxruntime/core/framework/external_data_loader_manager.cc index 91161b1d3dd4c..5d1a3fa1c573c 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.cc +++ b/onnxruntime/core/framework/external_data_loader_manager.cc @@ -2,7 +2,6 @@ // Licensed under the MIT License. #include "core/framework/external_data_loader_manager.h" -#include "core/framework/tensor.h" namespace onnxruntime { using namespace common; @@ -26,4 +25,57 @@ const IExternalDataLoader* ExternalDataLoaderManager::GetExternalDataLoader(cons return nullptr; } +const IExternalDataLoader* ExternalDataLoaderManager::GetExternalDataLoader( + const OrtMemoryInfo& target_memory_info, int32_t tensor_data_type) const { + for (auto& external_data_loader : external_data_loaders_) { + if (external_data_loader->CanLoad(target_memory_info) && + external_data_loader->SupportsDataType(tensor_data_type)) { + return external_data_loader.get(); + } + } + return nullptr; +} + +const IExternalDataLoader* ExternalDataLoaderManager::GetTensorCreator( + const OrtDevice& target_device, int32_t tensor_data_type) const { + for (const auto& external_data_loader : external_data_loaders_) { + if (external_data_loader->SupportsDataType(tensor_data_type) && + external_data_loader->CreatesTensorForDevice(target_device)) { + return external_data_loader.get(); + } + } + + return nullptr; +} + +Status ExternalDataLoaderManager::BeginLoad() const { + for (const auto& external_data_loader : external_data_loaders_) { + auto status = external_data_loader->BeginLoad(); + if (!status.IsOK()) { + AbortLoad(); + return status; + } + } + + return Status::OK(); +} + +Status ExternalDataLoaderManager::FinalizeLoad(const std::function& is_canceled) const { + for (const auto& external_data_loader : external_data_loaders_) { + auto status = external_data_loader->FinalizeLoad(is_canceled); + if (!status.IsOK()) { + AbortLoad(); + return status; + } + } + + return Status::OK(); +} + +void ExternalDataLoaderManager::AbortLoad() const noexcept { + for (const auto& external_data_loader : external_data_loaders_) { + external_data_loader->AbortLoad(); + } +} + } // namespace onnxruntime diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index c2bcd1c9034e7..5fcd6d525ebde 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.h +++ b/onnxruntime/core/framework/external_data_loader_manager.h @@ -22,6 +22,14 @@ class ExternalDataLoaderManager { // Release initialization-only loaders without invalidating SessionState references to this manager. void Clear() noexcept { external_data_loaders_.clear(); } + const IExternalDataLoader* GetExternalDataLoader( + const OrtMemoryInfo& target_memory_info, int32_t tensor_data_type) const; + const IExternalDataLoader* GetTensorCreator( + const OrtDevice& target_device, int32_t tensor_data_type) const; + common::Status BeginLoad() const; + common::Status FinalizeLoad(const std::function& is_canceled) const; + void AbortLoad() const noexcept; + private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExternalDataLoaderManager); diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index db988873539ad..92f2b432b447a 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -100,12 +100,20 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st const auto device = memory_info.device; if (utils::HasExternalData(tensor_proto)) { - auto external_data_loader = external_data_loader_mgr.GetExternalDataLoader(memory_info); + auto external_data_loader = + external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); if (external_data_loader) { - // if custom external data loader is used, always allocate memory on device - ORT_RETURN_IF_ERROR(AllocateTensor(memory_buffer, tensor, type, tensor_shape, use_device_allocator_for_initializers, alloc)); + if (external_data_loader->CreatesTensorForDevice(device)) { + ORT_RETURN_IF(memory_buffer != nullptr || alloc == nullptr, + "An external data loader that creates tensors requires a device allocator."); + tensor = Tensor{type, tensor_shape, nullptr, alloc}; + } else { + // if custom external data loader is used, always allocate memory on device + ORT_RETURN_IF_ERROR( + AllocateTensor(memory_buffer, tensor, type, tensor_shape, use_device_allocator_for_initializers, alloc)); + } ORT_RETURN_IF_ERROR(utils::LoadExtDataToTensorFromTensorProto(env, proto_path, tensor_proto, - *external_data_loader, tensor)); + *external_data_loader, alloc, tensor)); ORT_RETURN_IF_ERROR(utils::LoadPrepackedWeightsFromExternalData(env, proto_path, tensor_proto, prepacked_for_graph)); @@ -330,8 +338,14 @@ common::Status SaveInitializedTensors( // - Values that are external and mapped from disk. We let the OS manage the memory. // - we do not trace values that are in memory because they may be sitting on top of the user allocated // memory. - const bool trace_allocation = (exec_plan.GetLocation(ort_value_index) != default_cpu_device) || - !utils::HasExternalData(*tensor_proto); + const bool loader_creates_tensor = + utils::HasExternalData(*tensor_proto) && + !utils::HasExternalDataInMemory(*tensor_proto) && + external_data_loader_mgr.GetTensorCreator( + exec_plan.GetLocation(ort_value_index), tensor_proto->data_type()) != nullptr; + const bool trace_allocation = !loader_creates_tensor && + ((exec_plan.GetLocation(ort_value_index) != default_cpu_device) || + !utils::HasExternalData(*tensor_proto)); if (trace_allocation) { // can not trace string tensor, and they exist only on CPU @@ -350,6 +364,12 @@ common::Status SaveInitializedTensors( // do not trace string tensor continue; } + if (utils::HasExternalData(*entry.second) && + !utils::HasExternalDataInMemory(*entry.second) && + external_data_loader_mgr.GetTensorCreator( + exec_plan.GetLocation(entry.first), entry.second->data_type()) != nullptr) { + continue; + } ORT_RETURN_IF_ERROR(planner.Trace(entry.first, entry.second)); } @@ -375,6 +395,37 @@ common::Status SaveInitializedTensors( session_options.config_options.GetConfigOrDefault( kOrtSessionOptionsUseDeviceAllocatorForInitializers, "0") == "1"; + ORT_RETURN_IF_ERROR(external_data_loader_mgr.BeginLoad()); + bool external_data_load_finalized = false; + auto abort_external_data_load = gsl::finally([&]() { + if (!external_data_load_finalized) { + external_data_loader_mgr.AbortLoad(); + } + }); + + for (const auto& entry : id_to_initialized_tensor) { + if (user_supplied_initializer_ids.contains(entry.first) || + !utils::HasExternalData(*entry.second) || + utils::HasExternalDataInMemory(*entry.second)) { + continue; + } + + if (session_options.IsLoadCancellationFlagSet()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, MODEL_LOAD_CANCELED, + "Preparing session state weights is canceled due to user request."); + } + + const auto* tensor_creator = external_data_loader_mgr.GetTensorCreator( + exec_plan.GetLocation(entry.first), entry.second->data_type()); + if (tensor_creator != nullptr) { + ORT_RETURN_IF_ERROR( + utils::PrepareExtDataForTensorFromTensorProto(env, graph_loc, *entry.second, *tensor_creator)); + } + } + + ORT_RETURN_IF_ERROR(external_data_loader_mgr.FinalizeLoad( + [&session_options]() { return session_options.IsLoadCancellationFlagSet(); })); + // 3. create weight tensors based on weights buffer for (const auto& entry : id_to_initialized_tensor) { // We check for cancellation for every initializer since mapping from disk can be costly @@ -462,6 +513,7 @@ common::Status SaveInitializedTensors( #endif } + external_data_load_finalized = true; LOGS(logger, INFO) << "Done saving initialized tensors"; return common::Status::OK(); } diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 52a4cfe90a10b..723610d984629 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -1958,6 +1958,7 @@ Status LoadPrepackedWeightsFromExternalData(const Env& env, Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, const ONNX_NAMESPACE::TensorProto& tensor_proto, const IExternalDataLoader& ext_data_loader, + const AllocatorPtr& allocator, Tensor& tensor) { ORT_ENFORCE(HasExternalData(tensor_proto)); // Defense-in-depth path validation for callers reaching this function outside Graph::Resolve. @@ -1980,6 +1981,13 @@ Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem: ORT_RETURN_IF(external_data_file_path == onnxruntime::utils::kTensorProtoLittleEndianMemoryAddressTag || external_data_file_path == onnxruntime::utils::kTensorProtoNativeEndianMemoryAddressTag, "Memory address tag is not supported by custom external data loader."); +#if defined(ENABLE_D3D12_FILE_LOADING) + if (ext_data_loader.CreatesTensorForDevice(tensor.Location().device)) { + return ext_data_loader.LoadTensor(env, external_data_file_path, tensor_proto.name(), file_offset, + raw_data_safe_len, allocator, tensor); + } +#endif + #if defined(__wasm__) return ext_data_loader.LoadTensor(env, external_data_file_path, file_offset, raw_data_safe_len, tensor); #else @@ -1990,6 +1998,40 @@ Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem: #endif } +#if defined(ENABLE_D3D12_FILE_LOADING) +Status PrepareExtDataForTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, + const IExternalDataLoader& ext_data_loader) { + ORT_ENFORCE(HasExternalData(tensor_proto)); + + std::basic_string tensor_proto_dir; + if (!model_path.empty()) { + ORT_RETURN_IF_ERROR(GetDirNameFromFilePath(model_path, tensor_proto_dir)); + } + + std::basic_string external_data_file_path; + FileOffsetType file_offset; + SafeInt raw_data_safe_len = 0; + ORT_RETURN_IF_ERROR( + GetExternalDataInfo(tensor_proto, tensor_proto_dir, external_data_file_path, file_offset, raw_data_safe_len)); + ORT_RETURN_IF_ERROR(ValidateExternalFilePathForTensor(tensor_proto, model_path)); + + size_t tensor_byte_size = 0; + ORT_RETURN_IF_ERROR(GetSizeInBytesFromTensorProto<0>(tensor_proto, &tensor_byte_size)); + ORT_RETURN_IF(file_offset < 0 || raw_data_safe_len != tensor_byte_size, + "External initializer: ", tensor_proto.name(), " offset: ", file_offset, + " size to read: ", static_cast(raw_data_safe_len), + " does not match the tensor size: ", tensor_byte_size); + ORT_RETURN_IF(external_data_file_path == onnxruntime::utils::kTensorProtoLittleEndianMemoryAddressTag || + external_data_file_path == onnxruntime::utils::kTensorProtoNativeEndianMemoryAddressTag, + "Memory address tag is not supported by custom external data loader."); + + return ext_data_loader.PrepareTensor( + env, external_data_file_path, tensor_proto.name(), file_offset, + raw_data_safe_len); +} +#endif + #define CASE_PROTO(X, Y) \ case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \ ORT_RETURN_IF_ERROR( \ diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 19aff139cfad1..655dba628b183 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -325,6 +325,7 @@ common::Status GetExtDataFromTensorProto(const Env& env, const std::filesystem:: common::Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, const ONNX_NAMESPACE::TensorProto& tensor_proto, const IExternalDataLoader& ext_data_loader, + const AllocatorPtr& allocator, Tensor& tensor); // Load any saved pre-packed blobs referenced by an external TensorProto without loading its tensor data. @@ -332,6 +333,10 @@ common::Status LoadPrepackedWeightsFromExternalData(const Env& env, const std::f const ONNX_NAMESPACE::TensorProto& tensor_proto, PrepackedWeightsForGraph& prepacked_info); +common::Status PrepareExtDataForTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, + const IExternalDataLoader& ext_data_loader); + // Convert the AttributeProto from a Constant node into a TensorProto that can be used as an initializer // If AttributeProto contains a TensorProto, this tensor proto is converted as is including the case when the // the data location is external. i.e. it does not load the external data. diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc new file mode 100644 index 0000000000000..62c21f770ae82 --- /dev/null +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -0,0 +1,1067 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "core/common/common.h" + +namespace onnxruntime { +namespace windows { +namespace d3d12 { + +using Microsoft::WRL::ComPtr; + +namespace { + +constexpr uint64_t kBufferAlignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT; +constexpr DWORD kCancellationPollMilliseconds = 10; + +common::Status HResultError(const char* operation, HRESULT hr) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, operation, " failed with HRESULT ", + static_cast(hr)); +} + +common::Status Win32Error(const char* operation, DWORD error) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, operation, " failed, error code ", error, " - ", + std::system_category().message(error)); +} + +common::Status CancelledStatus() { + return ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 file buffer loading was canceled."); +} + +uint64_t AlignDown(uint64_t value, uint64_t alignment) noexcept { + return value - value % alignment; +} + +bool TryAlignUp(uint64_t value, uint64_t alignment, uint64_t& result) noexcept { + if (alignment == 0 || + value > std::numeric_limits::max() - (alignment - 1)) { + return false; + } + result = AlignDown(value + alignment - 1, alignment); + return true; +} + +D3D12_HEAP_PROPERTIES HeapProperties(D3D12_HEAP_TYPE type) noexcept { + D3D12_HEAP_PROPERTIES properties{}; + properties.Type = type; + properties.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN; + properties.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN; + properties.CreationNodeMask = 1; + properties.VisibleNodeMask = 1; + return properties; +} + +D3D12_RESOURCE_DESC BufferDescription( + uint64_t size, D3D12_RESOURCE_FLAGS flags) noexcept { + D3D12_RESOURCE_DESC description{}; + description.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; + description.Width = size; + description.Height = 1; + description.DepthOrArraySize = 1; + description.MipLevels = 1; + description.Format = DXGI_FORMAT_UNKNOWN; + description.SampleDesc.Count = 1; + description.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; + description.Flags = flags; + return description; +} + +} // namespace + +struct D3D12FileBufferLoader::Impl { + struct FileReadRegion { + uint64_t begin = 0; + uint64_t end = 0; + }; + + struct PreparedFile { + std::wstring path; + wil::unique_hfile file; + uint64_t size = 0; + uint64_t alignment = 0; + std::vector range_indices; + std::vector regions; + }; + + struct UploadSlot { + ComPtr resource; + ComPtr allocator; + ComPtr command_list; + wil::unique_handle read_event; + void* mapped = nullptr; + OVERLAPPED overlapped{}; + uint64_t file_offset = 0; + DWORD requested = 0; + uint64_t fence_value = 0; + bool read_active = false; + }; + + struct AllocationResult { + common::Status status; + Batch batch; + }; + + Impl(ID3D12Device* device, const Config& config) + : device_(device), config_(config) { + } + + ~Impl() { + for (auto& slot : slots_) { + if (slot.resource && slot.mapped != nullptr) { + slot.resource->Unmap(0, nullptr); + slot.mapped = nullptr; + } + } + } + + common::Status Initialize() { + if (device_ == nullptr) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "D3D12FileBufferLoader requires a non-null D3D12 device."); + } + if (config_.upload_slot_count < 2 || + config_.upload_slot_count > MAXIMUM_WAIT_OBJECTS) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "upload_slot_count must be between 2 and ", MAXIMUM_WAIT_OBJECTS, "."); + } + if (config_.upload_slot_size == 0 || + config_.upload_slot_size > std::numeric_limits::max() || + config_.upload_slot_size % kBufferAlignment != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "upload_slot_size must be a non-zero 64 KiB multiple no larger than DWORD_MAX."); + } + if (config_.max_destination_heap_size == 0 || + config_.max_destination_heap_size > 512ull * 1024ull * 1024ull || + config_.max_destination_heap_size % kBufferAlignment != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "max_destination_heap_size must be a non-zero 64 KiB multiple no larger than 512 MiB."); + } + + D3D12_COMMAND_QUEUE_DESC queue_description{}; + queue_description.Type = D3D12_COMMAND_LIST_TYPE_COPY; + HRESULT hr = device_->CreateCommandQueue( + &queue_description, IID_PPV_ARGS(©_queue_)); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateCommandQueue", hr); + } + + hr = device_->CreateFence( + 0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(©_fence_)); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateFence", hr); + } + + fence_event_.reset(CreateEventExW( + nullptr, nullptr, 0, EVENT_MODIFY_STATE | SYNCHRONIZE)); + if (!fence_event_) { + return Win32Error("CreateEventExW(copy fence)", GetLastError()); + } + + const auto upload_heap = HeapProperties(D3D12_HEAP_TYPE_UPLOAD); + const auto upload_description = + BufferDescription(config_.upload_slot_size, D3D12_RESOURCE_FLAG_NONE); + slots_.resize(config_.upload_slot_count); + for (auto& slot : slots_) { + hr = device_->CreateCommittedResource( + &upload_heap, D3D12_HEAP_FLAG_NONE, &upload_description, + D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, + IID_PPV_ARGS(&slot.resource)); + if (FAILED(hr)) { + return HResultError( + "ID3D12Device::CreateCommittedResource(upload)", hr); + } + + hr = slot.resource->Map(0, nullptr, &slot.mapped); + if (FAILED(hr)) { + return HResultError("ID3D12Resource::Map(upload)", hr); + } + + hr = device_->CreateCommandAllocator( + D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&slot.allocator)); + if (FAILED(hr)) { + return HResultError( + "ID3D12Device::CreateCommandAllocator", hr); + } + + hr = device_->CreateCommandList( + 0, D3D12_COMMAND_LIST_TYPE_COPY, slot.allocator.Get(), nullptr, + IID_PPV_ARGS(&slot.command_list)); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateCommandList", hr); + } + hr = slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close(initial)", hr); + } + + slot.read_event.reset(CreateEventExW( + nullptr, nullptr, CREATE_EVENT_MANUAL_RESET, + EVENT_MODIFY_STATE | SYNCHRONIZE)); + if (!slot.read_event) { + return Win32Error("CreateEventExW(read)", GetLastError()); + } + } + + return common::Status::OK(); + } + + common::Status Load( + const std::vector& ranges, + Batch& result, + const CancellationToken& cancellation) { + std::lock_guard lock(load_mutex_); + result.Clear(); + + if (ranges.empty()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "At least one file range is required."); + } + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + + std::vector sizes; + sizes.reserve(ranges.size()); + for (size_t index = 0; index < ranges.size(); ++index) { + const auto& range = ranges[index]; + if (range.path.empty()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "File range ", index, " has an empty path."); + } + if (range.length == 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "File range ", index, " has zero length."); + } + if (range.offset > + std::numeric_limits::max() - range.length) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "File range ", index, " overflows uint64_t."); + } + sizes.push_back(range.length); + } + + std::future allocation_future = std::async( + std::launch::async, [this, sizes = std::move(sizes)]() noexcept { + AllocationResult allocation; + try { + allocation.status = + AllocateDestinations(sizes, allocation.batch); + } catch (const std::exception& ex) { + allocation.status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "Destination allocation failed: ", ex.what()); + } catch (...) { + allocation.status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "Destination allocation failed with an unknown exception."); + } + return allocation; + }); + + std::vector files; + HANDLE active_file = INVALID_HANDLE_VALUE; + uint64_t last_submitted_fence = 0; + bool load_succeeded = false; + auto cleanup = gsl::finally([&]() noexcept { + if (!load_succeeded) { + DrainActiveReads(active_file); + WaitForFenceUncancelled(last_submitted_fence); + } + if (allocation_future.valid()) { + try { + (void)allocation_future.get(); + } catch (...) { + } + } + }); + + ORT_RETURN_IF_ERROR(PrepareFiles(ranges, cancellation, files)); + + uint64_t maximum_alignment = 1; + for (const auto& file : files) { + maximum_alignment = std::max(maximum_alignment, file.alignment); + } + ORT_RETURN_IF_ERROR(PrepareSlots(maximum_alignment, cancellation)); + + Batch loaded_batch; + bool allocation_ready = false; + auto ensure_allocation = [&]() -> common::Status { + if (allocation_ready) { + return common::Status::OK(); + } + AllocationResult allocation = allocation_future.get(); + if (!allocation.status.IsOK()) { + return allocation.status; + } + loaded_batch = std::move(allocation.batch); + allocation_ready = true; + return common::Status::OK(); + }; + + for (auto& file : files) { + active_file = file.file.get(); + for (const auto& region : file.regions) { + const auto status = ReadRegion( + file, region, ranges, loaded_batch, allocation_ready, + ensure_allocation, cancellation, last_submitted_fence); + if (!status.IsOK()) { + return status; + } + } + active_file = INVALID_HANDLE_VALUE; + } + + ORT_RETURN_IF_ERROR(ensure_allocation()); + ORT_RETURN_IF_ERROR( + WaitForFence(last_submitted_fence, cancellation)); + ORT_RETURN_IF_ERROR( + TransitionToCommon(loaded_batch, cancellation, last_submitted_fence)); + + result = std::move(loaded_batch); + load_succeeded = true; + return common::Status::OK(); + } + + common::Status AllocateDestinations( + const std::vector& sizes, + Batch& batch) { + struct Placement { + bool committed = false; + size_t heap_index = 0; + uint64_t offset = 0; + D3D12_RESOURCE_DESC description{}; + }; + + std::vector placements; + placements.reserve(sizes.size()); + std::vector heap_sizes; + + for (size_t index = 0; index < sizes.size(); ++index) { + uint64_t resource_size = 0; + if (!TryAlignUp(sizes[index], 16, resource_size)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "Destination ", index, " size overflows while aligning for Dawn."); + } + const auto description = BufferDescription( + resource_size, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS); + const auto allocation_info = + device_->GetResourceAllocationInfo(0, 1, &description); + if (allocation_info.SizeInBytes == + std::numeric_limits::max()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "GetResourceAllocationInfo failed for destination ", index, "."); + } + if (allocation_info.SizeInBytes > config_.max_destination_heap_size) { + placements.push_back({true, 0, 0, description}); + continue; + } + + if (heap_sizes.empty()) { + heap_sizes.push_back(0); + } + size_t heap_index = heap_sizes.size() - 1; + uint64_t offset = 0; + if (!TryAlignUp( + heap_sizes.back(), allocation_info.Alignment, offset)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Destination placement offset overflow."); + } + if (heap_sizes.back() != 0 && + offset > + config_.max_destination_heap_size - + allocation_info.SizeInBytes) { + heap_sizes.push_back(0); + heap_index = heap_sizes.size() - 1; + offset = 0; + } + + placements.push_back({false, heap_index, offset, description}); + heap_sizes[heap_index] = + offset + allocation_info.SizeInBytes; + } + + batch.heaps.resize(heap_sizes.size()); + for (size_t index = 0; index < heap_sizes.size(); ++index) { + uint64_t heap_size = 0; + if (!TryAlignUp( + heap_sizes[index], kBufferAlignment, heap_size) || + heap_size > config_.max_destination_heap_size) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Destination heap size is invalid."); + } + + D3D12_HEAP_DESC heap_description{}; + heap_description.SizeInBytes = heap_size; + heap_description.Properties = + HeapProperties(D3D12_HEAP_TYPE_DEFAULT); + heap_description.Flags = D3D12_HEAP_FLAG_ALLOW_ONLY_BUFFERS; + const HRESULT hr = device_->CreateHeap( + &heap_description, IID_PPV_ARGS(&batch.heaps[index])); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateHeap", hr); + } + } + + batch.buffers.resize(sizes.size()); + for (size_t index = 0; index < placements.size(); ++index) { + const auto& placement = placements[index]; + HRESULT hr = S_OK; + if (placement.committed) { + const auto heap_properties = HeapProperties(D3D12_HEAP_TYPE_DEFAULT); + hr = device_->CreateCommittedResource( + &heap_properties, D3D12_HEAP_FLAG_NONE, &placement.description, + D3D12_RESOURCE_STATE_COPY_DEST, nullptr, + IID_PPV_ARGS(&batch.buffers[index].resource)); + } else { + hr = device_->CreatePlacedResource( + batch.heaps[placement.heap_index].Get(), + placement.offset, &placement.description, + D3D12_RESOURCE_STATE_COPY_DEST, nullptr, + IID_PPV_ARGS(&batch.buffers[index].resource)); + } + if (FAILED(hr)) { + return HResultError( + placement.committed + ? "ID3D12Device::CreateCommittedResource(destination)" + : "ID3D12Device::CreatePlacedResource", + hr); + } + batch.buffers[index].size = sizes[index]; + } + + return common::Status::OK(); + } + + common::Status PrepareFiles( + const std::vector& ranges, + const CancellationToken& cancellation, + std::vector& files) { + for (size_t range_index = 0; + range_index < ranges.size(); ++range_index) { + const auto& range = ranges[range_index]; + auto file_it = std::find_if( + files.begin(), files.end(), + [&](const PreparedFile& file) { + return file.path == range.path; + }); + if (file_it == files.end()) { + PreparedFile file; + file.path = range.path; + files.push_back(std::move(file)); + file_it = std::prev(files.end()); + } + file_it->range_indices.push_back(range_index); + } + + for (auto& file : files) { + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + + file.file.reset(CreateFileW( + file.path.c_str(), GENERIC_READ, FILE_SHARE_READ, + nullptr, OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED | + FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN, + nullptr)); + if (!file.file) { + return Win32Error( + "CreateFileW(unbuffered)", GetLastError()); + } + + LARGE_INTEGER file_size{}; + if (!GetFileSizeEx(file.file.get(), &file_size)) { + return Win32Error("GetFileSizeEx", GetLastError()); + } + if (file_size.QuadPart < 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, "Source file has a negative size."); + } + file.size = static_cast(file_size.QuadPart); + + FILE_STORAGE_INFO storage_info{}; + if (!GetFileInformationByHandleEx( + file.file.get(), FileStorageInfo, + &storage_info, sizeof(storage_info))) { + return Win32Error( + "GetFileInformationByHandleEx(FileStorageInfo)", + GetLastError()); + } + file.alignment = std::max( + storage_info.LogicalBytesPerSector, + storage_info.PhysicalBytesPerSectorForPerformance); + if (file.alignment == 0 || + (file.alignment & (file.alignment - 1)) != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Storage alignment is not a non-zero power of two."); + } + if (config_.upload_slot_size % file.alignment != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "upload_slot_size does not satisfy source file alignment."); + } + + std::sort( + file.range_indices.begin(), file.range_indices.end(), + [&](size_t left, size_t right) { + return ranges[left].offset < ranges[right].offset; + }); + + for (size_t range_index : file.range_indices) { + const auto& range = ranges[range_index]; + if (range.offset > file.size || + range.length > file.size - range.offset) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "File range ", range_index, + " exceeds its source file."); + } + + uint64_t aligned_end = 0; + uint64_t aligned_file_size = 0; + if (!TryAlignUp( + range.offset + range.length, + file.alignment, aligned_end) || + !TryAlignUp( + file.size, file.alignment, + aligned_file_size)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "Aligned file range overflow."); + } + FileReadRegion region{ + AlignDown(range.offset, file.alignment), + std::min(aligned_end, aligned_file_size)}; + if (!file.regions.empty() && + region.begin <= file.regions.back().end) { + file.regions.back().end = + std::max(file.regions.back().end, region.end); + } else { + file.regions.push_back(region); + } + } + } + + return common::Status::OK(); + } + + common::Status PrepareSlots( + uint64_t alignment, + const CancellationToken& cancellation) { + for (auto& slot : slots_) { + ORT_RETURN_IF_ERROR( + WaitForFence(slot.fence_value, cancellation)); + if (reinterpret_cast(slot.mapped) % + alignment != + 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Mapped upload buffer does not satisfy source file alignment."); + } + slot.overlapped = {}; + slot.file_offset = 0; + slot.requested = 0; + slot.fence_value = 0; + slot.read_active = false; + } + return common::Status::OK(); + } + + common::Status IssueRead( + HANDLE file, + UploadSlot& slot, + uint64_t offset, + DWORD size) { + if (!ResetEvent(slot.read_event.get())) { + return Win32Error("ResetEvent(read)", GetLastError()); + } + slot.overlapped = {}; + slot.overlapped.Offset = static_cast(offset); + slot.overlapped.OffsetHigh = + static_cast(offset >> 32); + slot.overlapped.hEvent = slot.read_event.get(); + slot.file_offset = offset; + slot.requested = size; + slot.read_active = true; + if (!ReadFile( + file, slot.mapped, size, nullptr, + &slot.overlapped)) { + const DWORD error = GetLastError(); + if (error != ERROR_IO_PENDING) { + slot.read_active = false; + return Win32Error("ReadFile(unbuffered)", error); + } + } + return common::Status::OK(); + } + + common::Status CompleteRead( + HANDLE file, + uint64_t file_size, + UploadSlot& slot, + DWORD& bytes_read) { + if (!GetOverlappedResult( + file, &slot.overlapped, + &bytes_read, FALSE)) { + return Win32Error( + "GetOverlappedResult", GetLastError()); + } + slot.read_active = false; + const uint64_t required_bytes = + std::min( + slot.requested, + file_size - slot.file_offset); + if (bytes_read < required_bytes) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Unbuffered read ended before the required file range."); + } + return common::Status::OK(); + } + + template + common::Status ReadRegion( + PreparedFile& file, + const FileReadRegion& region, + const std::vector& ranges, + Batch& batch, + bool& allocation_ready, + EnsureAllocationFn& ensure_allocation, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence) { + uint64_t next_offset = region.begin; + + auto issue_available = [&]() -> common::Status { + for (auto& slot : slots_) { + if (next_offset >= region.end) { + break; + } + if (slot.read_active) { + continue; + } + if (slot.fence_value != 0 && + copy_fence_->GetCompletedValue() < + slot.fence_value) { + continue; + } + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + const DWORD bytes = static_cast( + std::min( + config_.upload_slot_size, + region.end - next_offset)); + ORT_RETURN_IF_ERROR( + IssueRead( + file.file.get(), slot, + next_offset, bytes)); + next_offset += bytes; + } + return common::Status::OK(); + }; + + ORT_RETURN_IF_ERROR(issue_available()); + while (true) { + std::array handles{}; + std::array slot_indices{}; + DWORD handle_count = 0; + for (size_t index = 0; + index < slots_.size(); ++index) { + if (slots_[index].read_active) { + handles[handle_count] = + slots_[index].read_event.get(); + slot_indices[handle_count] = index; + ++handle_count; + } + } + + if (handle_count == 0) { + if (next_offset >= region.end) { + break; + } + uint64_t next_fence = + std::numeric_limits::max(); + for (const auto& slot : slots_) { + if (!slot.read_active && + slot.fence_value != 0) { + next_fence = + std::min(next_fence, slot.fence_value); + } + } + if (next_fence == + std::numeric_limits::max()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Event-driven read scheduler made no progress."); + } + ORT_RETURN_IF_ERROR( + WaitForFence(next_fence, cancellation)); + ORT_RETURN_IF_ERROR(issue_available()); + continue; + } + + const DWORD wait_result = WaitForMultipleObjects( + handle_count, handles.data(), FALSE, + kCancellationPollMilliseconds); + if (wait_result == WAIT_TIMEOUT) { + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + continue; + } + if (wait_result < WAIT_OBJECT_0 || + wait_result >= WAIT_OBJECT_0 + handle_count) { + return Win32Error( + "WaitForMultipleObjects(read)", + GetLastError()); + } + + auto& slot = + slots_[slot_indices[wait_result - WAIT_OBJECT_0]]; + DWORD bytes_read = 0; + ORT_RETURN_IF_ERROR( + CompleteRead( + file.file.get(), file.size, + slot, bytes_read)); + + if (!allocation_ready) { + ORT_RETURN_IF_ERROR(ensure_allocation()); + } + ORT_RETURN_IF_ERROR( + SubmitCopies( + slot, slot.file_offset, bytes_read, + file.range_indices, ranges, batch, + last_submitted_fence)); + ORT_RETURN_IF_ERROR(issue_available()); + } + + return common::Status::OK(); + } + + common::Status SubmitCopies( + UploadSlot& slot, + uint64_t chunk_begin, + uint64_t bytes, + const std::vector& range_indices, + const std::vector& ranges, + Batch& batch, + uint64_t& last_submitted_fence) { + HRESULT hr = slot.allocator->Reset(); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandAllocator::Reset", hr); + } + hr = slot.command_list->Reset( + slot.allocator.Get(), nullptr); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Reset", hr); + } + + const uint64_t chunk_end = chunk_begin + bytes; + for (size_t range_index : range_indices) { + const auto& range = ranges[range_index]; + const uint64_t range_end = + range.offset + range.length; + const uint64_t copy_begin = + std::max(chunk_begin, range.offset); + const uint64_t copy_end = + std::min(chunk_end, range_end); + if (copy_begin < copy_end) { + slot.command_list->CopyBufferRegion( + batch.buffers[range_index].resource.Get(), + copy_begin - range.offset, + slot.resource.Get(), + copy_begin - chunk_begin, + copy_end - copy_begin); + } + + const uint64_t aligned_size = + (range.length + 15) & ~uint64_t{15}; + for (uint64_t padding_offset = range.length; + padding_offset < aligned_size;) { + const uint64_t source_relative_offset = + (padding_offset - range.length) % range.length; + const uint64_t source_offset = + range.offset + source_relative_offset; + if (source_offset >= chunk_begin && + source_offset < chunk_end) { + const uint64_t copy_size = std::min( + {aligned_size - padding_offset, + range.length - source_relative_offset, + chunk_end - source_offset}); + slot.command_list->CopyBufferRegion( + batch.buffers[range_index].resource.Get(), + padding_offset, + slot.resource.Get(), + source_offset - chunk_begin, + copy_size); + padding_offset += copy_size; + } else { + ++padding_offset; + } + } + } + + hr = slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close", hr); + } + ID3D12CommandList* command_lists[] = { + slot.command_list.Get()}; + copy_queue_->ExecuteCommandLists(1, command_lists); + + slot.fence_value = ++next_fence_value_; + last_submitted_fence = slot.fence_value; + hr = copy_queue_->Signal( + copy_fence_.Get(), slot.fence_value); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandQueue::Signal", hr); + } + return common::Status::OK(); + } + + common::Status TransitionToCommon( + Batch& batch, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence) { + if (batch.buffers.empty()) { + return common::Status::OK(); + } + + auto& slot = slots_.front(); + HRESULT hr = slot.allocator->Reset(); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandAllocator::Reset(transition)", hr); + } + hr = slot.command_list->Reset( + slot.allocator.Get(), nullptr); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Reset(transition)", hr); + } + + std::vector barriers; + barriers.reserve(batch.buffers.size()); + for (const auto& buffer : batch.buffers) { + D3D12_RESOURCE_BARRIER barrier{}; + barrier.Type = + D3D12_RESOURCE_BARRIER_TYPE_TRANSITION; + barrier.Transition.pResource = + buffer.resource.Get(); + barrier.Transition.Subresource = + D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES; + barrier.Transition.StateBefore = + D3D12_RESOURCE_STATE_COPY_DEST; + barrier.Transition.StateAfter = + D3D12_RESOURCE_STATE_COMMON; + barriers.push_back(barrier); + } + slot.command_list->ResourceBarrier( + static_cast(barriers.size()), + barriers.data()); + + hr = slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close(transition)", hr); + } + ID3D12CommandList* command_lists[] = { + slot.command_list.Get()}; + copy_queue_->ExecuteCommandLists(1, command_lists); + + last_submitted_fence = ++next_fence_value_; + hr = copy_queue_->Signal( + copy_fence_.Get(), last_submitted_fence); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandQueue::Signal(transition)", hr); + } + return WaitForFence( + last_submitted_fence, cancellation); + } + + common::Status WaitForFence( + uint64_t value, + const CancellationToken& cancellation) { + if (value == 0 || + copy_fence_->GetCompletedValue() >= value) { + return common::Status::OK(); + } + const HRESULT hr = + copy_fence_->SetEventOnCompletion( + value, fence_event_.get()); + if (FAILED(hr)) { + return HResultError( + "ID3D12Fence::SetEventOnCompletion", hr); + } + + while (true) { + const DWORD wait_result = WaitForSingleObject( + fence_event_.get(), + kCancellationPollMilliseconds); + if (wait_result == WAIT_OBJECT_0) { + return common::Status::OK(); + } + if (wait_result != WAIT_TIMEOUT) { + return Win32Error( + "WaitForSingleObject(copy fence)", + GetLastError()); + } + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + } + } + + void WaitForFenceUncancelled(uint64_t value) noexcept { + if (value == 0 || + copy_fence_->GetCompletedValue() >= value) { + return; + } + if (SUCCEEDED(copy_fence_->SetEventOnCompletion( + value, fence_event_.get())) && + WaitForSingleObject( + fence_event_.get(), INFINITE) == + WAIT_OBJECT_0) { + return; + } + + while (true) { + const uint64_t completed = + copy_fence_->GetCompletedValue(); + if (completed >= value || + completed == std::numeric_limits::max()) { + return; + } + Sleep(1); + } + } + + void DrainActiveReads(HANDLE file) noexcept { + if (file == nullptr || + file == INVALID_HANDLE_VALUE) { + return; + } + for (auto& slot : slots_) { + if (!slot.read_active) { + continue; + } + (void)CancelIoEx(file, &slot.overlapped); + } + for (auto& slot : slots_) { + if (!slot.read_active) { + continue; + } + DWORD bytes_read = 0; + (void)GetOverlappedResult( + file, &slot.overlapped, + &bytes_read, TRUE); + slot.read_active = false; + } + } + + ComPtr device_; + Config config_; + ComPtr copy_queue_; + ComPtr copy_fence_; + wil::unique_handle fence_event_; + std::vector slots_; + uint64_t next_fence_value_ = 0; + std::mutex load_mutex_; +}; + +D3D12FileBufferLoader::D3D12FileBufferLoader( + std::unique_ptr impl) noexcept + : impl_(std::move(impl)) { +} + +D3D12FileBufferLoader::~D3D12FileBufferLoader() = default; + +common::Status D3D12FileBufferLoader::Create( + ID3D12Device* device, + std::unique_ptr& loader, + const Config& config) noexcept { + loader.reset(); + try { + auto impl = std::make_unique(device, config); + ORT_RETURN_IF_ERROR(impl->Initialize()); + loader.reset( + new D3D12FileBufferLoader(std::move(impl))); + return common::Status::OK(); + } catch (const std::exception& ex) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "D3D12FileBufferLoader::Create failed: ", + ex.what()); + } catch (...) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "D3D12FileBufferLoader::Create failed with an unknown exception."); + } +} + +common::Status D3D12FileBufferLoader::Load( + const std::vector& ranges, + Batch& result, + const CancellationToken& cancellation) noexcept { + if (!impl_) { + result.Clear(); + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "D3D12FileBufferLoader is not initialized."); + } + try { + return impl_->Load(ranges, result, cancellation); + } catch (const std::exception& ex) { + result.Clear(); + return ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "D3D12FileBufferLoader::Load failed: ", + ex.what()); + } catch (...) { + result.Clear(); + return ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "D3D12FileBufferLoader::Load failed with an unknown exception."); + } +} + +} // namespace d3d12 +} // namespace windows +} // namespace onnxruntime diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h new file mode 100644 index 0000000000000..13d4cfb208b5b --- /dev/null +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h @@ -0,0 +1,88 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include "core/common/status.h" + +namespace onnxruntime { +namespace windows { +namespace d3d12 { + +class D3D12FileBufferLoader { + public: + struct Config { + uint32_t upload_slot_count = 12; + uint64_t upload_slot_size = 16ull * 1024ull * 1024ull; + uint64_t max_destination_heap_size = 512ull * 1024ull * 1024ull; + }; + + struct FileRange { + std::wstring path; + uint64_t offset = 0; + uint64_t length = 0; + }; + + struct CancellationToken { + using IsCancellationRequestedFn = bool (*)(void* state) noexcept; + + IsCancellationRequestedFn is_cancellation_requested = nullptr; + void* state = nullptr; + + bool IsCancellationRequested() const noexcept { + return is_cancellation_requested != nullptr && + is_cancellation_requested(state); + } + }; + + struct Buffer { + Microsoft::WRL::ComPtr resource; + uint64_t size = 0; + }; + + struct Batch { + std::vector> heaps; + std::vector buffers; + + void Clear() noexcept { + buffers.clear(); + heaps.clear(); + } + }; + + static common::Status Create( + ID3D12Device* device, + std::unique_ptr& loader, + const Config& config = {}) noexcept; + + ~D3D12FileBufferLoader(); + + D3D12FileBufferLoader(const D3D12FileBufferLoader&) = delete; + D3D12FileBufferLoader& operator=(const D3D12FileBufferLoader&) = delete; + + // The returned buffers correspond one-to-one with ranges and preserve input order. + // Each buffer is in D3D12_RESOURCE_STATE_COMMON when this method succeeds. + common::Status Load( + const std::vector& ranges, + Batch& result, + const CancellationToken& cancellation = {}) noexcept; + + private: + struct Impl; + + explicit D3D12FileBufferLoader(std::unique_ptr impl) noexcept; + + std::unique_ptr impl_; +}; + +} // namespace d3d12 +} // namespace windows +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/js/external_data_loader.cc b/onnxruntime/core/providers/js/external_data_loader.cc index 193b373cf3696..f5336d1754bad 100644 --- a/onnxruntime/core/providers/js/external_data_loader.cc +++ b/onnxruntime/core/providers/js/external_data_loader.cc @@ -20,8 +20,10 @@ bool ExternalDataLoader::CanLoad(const OrtMemoryInfo& target_memory_info) const common::Status ExternalDataLoader::LoadTensor(const Env& env, const std::filesystem::path& data_file_path, + std::string_view /*tensor_name*/, FileOffsetType data_offset, SafeInt data_length, + const std::shared_ptr& /*allocator*/, Tensor& tensor) const { ExternalDataLoadType load_type; if (tensor.Location().device.Type() == OrtDevice::CPU) { diff --git a/onnxruntime/core/providers/js/external_data_loader.h b/onnxruntime/core/providers/js/external_data_loader.h index 5f35ed62bbcc1..89fddb3d056ac 100644 --- a/onnxruntime/core/providers/js/external_data_loader.h +++ b/onnxruntime/core/providers/js/external_data_loader.h @@ -17,8 +17,10 @@ class ExternalDataLoader : public IExternalDataLoader { common::Status LoadTensor(const Env& env, const std::filesystem::path& data_file_path, + std::string_view tensor_name, FileOffsetType data_offset, SafeInt data_length, + const std::shared_ptr& allocator, Tensor& tensor) const override; }; diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index 46d81402cf8b7..8bfed74111d97 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -26,7 +26,7 @@ GpuBufferAllocator::GpuBufferAllocator( buffer_manager_getter_{std::move(buffer_manager_getter)}, recording_getter_{std::move(recording_getter)}, should_submit_zero_initialize_{std::move(should_submit_zero_initialize)}, - mapped_at_creation_{is_read_only_allocator && buffer_manager_getter_().SupportsUMA()}, + is_read_only_allocator_{is_read_only_allocator}, initialize_to_zero_{!is_read_only_allocator} { } @@ -49,11 +49,14 @@ void* GpuBufferAllocator::Allocate(size_t size, CommandRecordingState& recording return nullptr; } - wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite - : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; + const auto& buffer_manager = buffer_manager_getter_(); + if (!mapped_at_creation_.has_value()) { + mapped_at_creation_ = is_read_only_allocator_ && buffer_manager.SupportsUMA(); + } + wgpu::BufferUsage usage = *mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite + : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; - auto buffer = buffer_manager_getter_().Create(recording, size, usage, initialize_to_zero_, - submit_zero_initialize); + auto buffer = buffer_manager.Create(recording, size, usage, initialize_to_zero_, submit_zero_initialize); num_allocs_.fetch_add(1, std::memory_order_relaxed); return buffer; } diff --git a/onnxruntime/core/providers/webgpu/allocator.h b/onnxruntime/core/providers/webgpu/allocator.h index 04bb783d434eb..811daa3ac8d31 100644 --- a/onnxruntime/core/providers/webgpu/allocator.h +++ b/onnxruntime/core/providers/webgpu/allocator.h @@ -5,6 +5,7 @@ #include #include +#include #include "core/common/narrow.h" #include "core/framework/allocator.h" @@ -57,7 +58,8 @@ class GpuBufferAllocator : public IAllocator { std::function buffer_manager_getter_; std::function recording_getter_; std::function should_submit_zero_initialize_; - bool mapped_at_creation_; + bool is_read_only_allocator_; + std::optional mapped_at_creation_; // Cached writable buffers are cleared explicitly by BufferManager::Create. Fresh buffers rely on Dawn's // "lazy_clear_resource_on_first_use" toggle, which is enabled by WebGpuContext. bool initialize_to_zero_; diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc new file mode 100644 index 0000000000000..a9c9c9cdb783d --- /dev/null +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -0,0 +1,916 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/common/logging/logging.h" +#include "core/framework/tensor.h" +#include "core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h" +#include "core/providers/webgpu/allocator.h" +#include "core/providers/webgpu/buffer_manager.h" +#include "core/providers/webgpu/d3d12_external_data_loader.h" +#include "core/providers/webgpu/webgpu_context.h" + +namespace onnxruntime { +namespace webgpu { + +using Microsoft::WRL::ComPtr; +using D3D12FileBufferLoader = windows::d3d12::D3D12FileBufferLoader; + +namespace { + +using Clock = std::chrono::steady_clock; + +struct SharedBufferMemoryD3D12ResourceDescriptor : wgpu::ChainedStruct { + SharedBufferMemoryD3D12ResourceDescriptor() { + sType = static_cast( + WGPUSType_SharedBufferMemoryD3D12ResourceDescriptor); + } + + ComPtr resource; +}; + +double Milliseconds(Clock::duration duration) { + return std::chrono::duration(duration).count(); +} + +struct TensorKey { + std::filesystem::path path; + std::string name; + uint64_t offset; + size_t length; + + bool operator<(const TensorKey& other) const { + if (path != other.path) { + return path.native() < other.path.native(); + } + if (name != other.name) { + return name < other.name; + } + if (offset != other.offset) { + return offset < other.offset; + } + return length < other.length; + } +}; + +struct PreparedTensor { + TensorKey key; + ComPtr resource; + std::shared_ptr heap_owner; + wgpu::SharedBufferMemory memory; + wgpu::Buffer buffer; + bool access_started = false; + bool claimed = false; +}; + +struct D3D12AcceleratedBatch { + struct FileInfo { + std::filesystem::path canonical_path; + size_t length; + }; + + std::vector tensors; + std::map tensors_by_key; + std::map files; + size_t total_bytes = 0; + size_t load_range_count = 0; + bool finalized = false; +}; + +struct D3D12AcceleratedLoadMetrics { + double load_ms = 0.0; + size_t loaded_ranges = 0; +}; + +struct CancellationState { + const std::function* external = nullptr; + const std::atomic* abort_requested = nullptr; +}; + +bool IsCancellationRequested(void* opaque) noexcept { + const auto& state = *static_cast(opaque); + if (state.abort_requested != nullptr && + state.abort_requested->load(std::memory_order_relaxed)) { + return true; + } + if (state.external == nullptr || !*state.external) { + return false; + } + ORT_TRY { + return (*state.external)(); + } + ORT_CATCH(...) { + return true; + } +} + +common::Status LoadBatchToD3D12( + D3D12AcceleratedBatch& batch, + D3D12FileBufferLoader& loader, + const std::function& is_canceled, + const std::atomic& abort_requested, + D3D12AcceleratedLoadMetrics& metrics) { + if (abort_requested.load(std::memory_order_relaxed) || + (is_canceled && is_canceled())) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled."); + } + + std::vector ranges; + std::vector tensor_indices; + ranges.reserve(batch.load_range_count); + tensor_indices.reserve(batch.load_range_count); + for (size_t index = 0; index < batch.tensors.size(); ++index) { + const auto& tensor = batch.tensors[index]; + if (tensor.key.length == 0) { + continue; + } + ranges.push_back( + {tensor.key.path.native(), + tensor.key.offset, + static_cast(tensor.key.length)}); + tensor_indices.push_back(index); + } + + ORT_RETURN_IF_NOT( + ranges.size() == batch.load_range_count, + "D3D12 accelerated initializer range count is inconsistent."); + if (ranges.empty()) { + return common::Status::OK(); + } + + auto loaded = std::make_shared(); + const CancellationState cancellation_state{ + &is_canceled, &abort_requested}; + const D3D12FileBufferLoader::CancellationToken cancellation{ + IsCancellationRequested, + const_cast(&cancellation_state)}; + + const auto load_start = Clock::now(); + ORT_RETURN_IF_ERROR(loader.Load(ranges, *loaded, cancellation)); + const auto load_end = Clock::now(); + ORT_RETURN_IF_NOT( + loaded->buffers.size() == tensor_indices.size(), + "D3D12 file loader returned an unexpected buffer count."); + + for (size_t index = 0; index < tensor_indices.size(); ++index) { + auto& tensor = batch.tensors[tensor_indices[index]]; + tensor.resource = loaded->buffers[index].resource; + tensor.heap_owner = loaded; + } + + metrics.load_ms = Milliseconds(load_end - load_start); + metrics.loaded_ranges = ranges.size(); + return common::Status::OK(); +} + +TensorKey MakeKey(const std::filesystem::path& path, + std::string_view name, + uint64_t offset, + size_t length) { + return {path, std::string{name}, offset, length}; +} + +common::Status PrepareTensorForBatch( + D3D12AcceleratedBatch& batch, + const Env& env, + const std::filesystem::path& data_file_path, + std::string_view tensor_name, + FileOffsetType data_offset, + SafeInt data_length) { + ORT_RETURN_IF( + data_offset < 0, + "D3D12 accelerated initializer \"", tensor_name, + "\" has a negative file offset."); + + const size_t length = static_cast(data_length); + const uint64_t offset = static_cast(data_offset); + ORT_RETURN_IF( + offset > std::numeric_limits::max() - + static_cast(length), + "D3D12 accelerated initializer \"", tensor_name, + "\" file range overflows."); + + const auto file_key = data_file_path.lexically_normal(); + auto file_iterator = batch.files.find(file_key); + if (file_iterator == batch.files.end()) { + size_t file_length = 0; + ORT_RETURN_IF_ERROR( + env.GetFileLength(data_file_path.c_str(), file_length)); + std::error_code path_error; + const auto normalized_path = + std::filesystem::weakly_canonical( + data_file_path, path_error); + ORT_RETURN_IF( + path_error, + "Failed to canonicalize D3D12 data file \"", + data_file_path.string(), "\": ", + path_error.message()); + file_iterator = + batch.files + .emplace( + file_key, + D3D12AcceleratedBatch::FileInfo{ + normalized_path, file_length}) + .first; + } + + const size_t file_length = file_iterator->second.length; + ORT_RETURN_IF( + offset > static_cast(file_length) || + static_cast(length) > + static_cast(file_length) - offset, + "D3D12 accelerated initializer \"", tensor_name, + "\" range [", offset, ", ", offset + length, + ") exceeds file size ", file_length, " for \"", + data_file_path.string(), "\"."); + + TensorKey key = MakeKey( + file_iterator->second.canonical_path, + tensor_name, offset, length); + ORT_RETURN_IF( + batch.tensors_by_key.find(key) != + batch.tensors_by_key.end(), + "Duplicate D3D12 accelerated initializer: ", + tensor_name); + ORT_RETURN_IF( + batch.total_bytes > + std::numeric_limits::max() - length, + "D3D12 accelerated initializer byte count overflow."); + + const size_t index = batch.tensors.size(); + batch.tensors.push_back({std::move(key)}); + batch.tensors_by_key.emplace( + batch.tensors.back().key, index); + batch.total_bytes += length; + if (length != 0) { + ++batch.load_range_count; + } + return common::Status::OK(); +} + +struct ImportedAllocation { + ComPtr resource; + std::shared_ptr heap_owner; + wgpu::SharedBufferMemory memory; + wgpu::Buffer buffer; + bool access_started = false; +}; + +void EndAccessNoThrow( + ImportedAllocation& allocation) noexcept { + if (!allocation.access_started || + !allocation.memory || !allocation.buffer) { + return; + } + + ORT_TRY { + wgpu::SharedBufferMemoryEndAccessState end_state{}; + allocation.memory.EndAccess( + allocation.buffer, &end_state); + allocation.access_started = false; + } + ORT_CATCH(...) { + // Cleanup paths, including allocator Free, must not propagate Dawn failures. + } +} + +} // namespace + +common::Status CheckD3D12AcceleratedExternalWeightsSupport( + WebGpuContext& context) { + ORT_RETURN_IF_NOT( + context.RequestedBackendType() == + wgpu::BackendType::D3D12, + "D3D12 accelerated external weights require the Dawn D3D12 backend."); + ORT_RETURN_IF_NOT( + context.D3D12SharedResourceFeaturesAvailable(), + "D3D12 accelerated external weights require Dawn " + "SharedBufferMemoryD3D12Resource and SharedFenceDXGISharedHandle features."); + + ORT_RETURN_IF_NOT( + context.WeightLoadingD3D12Device() != nullptr, + "Failed to create a D3D12 device for accelerated file loading."); + return common::Status::OK(); +} + +common::Status ResolveWeightLoadAccelerationMode( + WeightLoadAccelerationMode mode, + const common::Status& support_status, + bool& enabled) { + enabled = false; + if (!IsWeightLoadAccelerationEnabled(mode)) { + return common::Status::OK(); + } + if (support_status.IsOK()) { + enabled = true; + return common::Status::OK(); + } + if (IsWeightLoadAccelerationRequired(mode)) { + return support_status; + } + return common::Status::OK(); +} + +struct D3D12AcceleratedInitializerState::Impl { + std::mutex mutex; + IAllocator* allocator = nullptr; + std::unordered_map< + WGPUBuffer, + std::unique_ptr> + imported_allocations; +}; + +D3D12AcceleratedInitializerState:: + D3D12AcceleratedInitializerState() = default; + +D3D12AcceleratedInitializerState:: + ~D3D12AcceleratedInitializerState() { + if (!impl_) { + return; + } + + std::vector> + allocations; + { + std::lock_guard lock{impl_->mutex}; + allocations.reserve( + impl_->imported_allocations.size()); + for (auto& entry : impl_->imported_allocations) { + allocations.push_back(std::move(entry.second)); + } + impl_->imported_allocations.clear(); + } + for (auto& allocation : allocations) { + EndAccessNoThrow(*allocation); + } +} + +class D3D12AcceleratedWebGpuAllocator final + : public IAllocator { + public: + D3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + std::function + recording_getter, + std::shared_ptr + state) + : IAllocator( + OrtMemoryInfo( + WEBGPU_BUFFER, + OrtAllocatorType::OrtReadOnlyAllocator, + WebGpuDevice, + OrtMemTypeDefault)), + context_{context}, + recording_getter_{ + std::move(recording_getter)}, + state_{std::move(state)} { + } + + void* Alloc(size_t size) override { + if (size == 0) { + return nullptr; + } + + const auto usage = + wgpu::BufferUsage::Storage | + wgpu::BufferUsage::CopySrc | + wgpu::BufferUsage::CopyDst | + wgpu::BufferUsage::Indirect; + auto& recording = recording_getter_(); + std::lock_guard lock{ + recording.mutex}; + return context_.InitializerBufferManager().Create( + recording, size, usage); + } + + void Free(void* p) override { + if (p == nullptr) { + return; + } + + auto& recording = recording_getter_(); + std::lock_guard + recording_lock{recording.mutex}; + std::unique_ptr imported; + { + std::lock_guard lock{ + state_->impl_->mutex}; + const auto iterator = + state_->impl_->imported_allocations.find( + static_cast(p)); + if (iterator != + state_->impl_->imported_allocations.end()) { + imported = std::move(iterator->second); + state_->impl_->imported_allocations.erase( + iterator); + } + } + + if (imported) { + if (recording.has_unsubmitted_work) { + std::shared_ptr deferred{ + imported.release(), + [](ImportedAllocation* allocation) { + EndAccessNoThrow(*allocation); + delete allocation; + }}; + recording.pending_release_callbacks.emplace_back( + [deferred = std::move(deferred)]() {}); + } else { + EndAccessNoThrow(*imported); + } + return; + } + + context_.InitializerBufferManager().Release( + recording, static_cast(p)); + } + + private: + WebGpuContext& context_; + std::function + recording_getter_; + std::shared_ptr + state_; +}; + +AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + std::function + recording_getter, + std::shared_ptr& + out_state) { + auto state = + std::shared_ptr( + new D3D12AcceleratedInitializerState()); + state->impl_ = + std::make_unique< + D3D12AcceleratedInitializerState::Impl>(); + auto allocator = + std::make_shared< + D3D12AcceleratedWebGpuAllocator>( + context, std::move(recording_getter), state); + state->impl_->allocator = allocator.get(); + out_state = state; + return allocator; +} + +struct D3D12AcceleratedExternalDataLoader::Impl { + Impl( + WebGpuContext& context_in, + std::shared_ptr + state_in, + WeightLoadAccelerationMode mode_in) + : context{context_in}, + state{std::move(state_in)}, + mode{mode_in} { + } + + void ResolveSupport() const { + std::call_once(support_once, [this]() { + common::Status support_status = + common::Status::OK(); + if (IsWeightLoadAccelerationEnabled(mode)) { + support_status = + CheckD3D12AcceleratedExternalWeightsSupport( + context); + } + resolved_status = + ResolveWeightLoadAccelerationMode( + mode, support_status, enabled); + if (resolved_status.IsOK() && !enabled && + !IsWeightLoadAccelerationRequired(mode)) { + LOGS_DEFAULT(WARNING) + << "D3D12 accelerated external weights are unavailable; " + "using the ordinary WebGPU initializer loading path. " + "Reason: " + << support_status.ErrorMessage(); + } + }); + } + + common::Status EnsureFileLoader() const { + if (file_loader) { + return common::Status::OK(); + } + return D3D12FileBufferLoader::Create( + context.WeightLoadingD3D12Device(), file_loader); + } + + WebGpuContext& context; + std::shared_ptr + state; + WeightLoadAccelerationMode mode; + mutable std::once_flag support_once; + mutable common::Status resolved_status; + mutable bool enabled = false; + mutable std::atomic abort_requested{false}; + mutable std::unique_ptr + file_loader; + mutable std::unique_ptr batch; +}; + +D3D12AcceleratedExternalDataLoader:: + D3D12AcceleratedExternalDataLoader( + WebGpuContext& context, + std::shared_ptr< + D3D12AcceleratedInitializerState> + state, + WeightLoadAccelerationMode mode) + : impl_{std::make_unique( + context, std::move(state), mode)} { + ORT_ENFORCE( + impl_->state != nullptr && + impl_->state->impl_ != nullptr, + "D3D12 accelerated allocator state is required."); +} + +D3D12AcceleratedExternalDataLoader:: + ~D3D12AcceleratedExternalDataLoader() { + AbortLoad(); +} + +bool D3D12AcceleratedExternalDataLoader::CanLoad( + const OrtMemoryInfo& target_memory_info) const { + impl_->ResolveSupport(); + return target_memory_info.device == WebGpuDevice && + target_memory_info.name == WEBGPU_BUFFER && + impl_->enabled; +} + +bool D3D12AcceleratedExternalDataLoader:: + SupportsDataType( + int32_t tensor_data_type) const { + return tensor_data_type != + ONNX_NAMESPACE::TensorProto_DataType_BOOL; +} + +bool D3D12AcceleratedExternalDataLoader:: + CreatesTensorForDevice( + const OrtDevice& target_device) const { + impl_->ResolveSupport(); + return target_device == WebGpuDevice && + (impl_->enabled || + IsWeightLoadAccelerationRequired(impl_->mode)); +} + +common::Status +D3D12AcceleratedExternalDataLoader::BeginLoad() const { + impl_->batch.reset(); + impl_->ResolveSupport(); + ORT_RETURN_IF_ERROR(impl_->resolved_status); + if (!impl_->enabled) { + return common::Status::OK(); + } + impl_->abort_requested.store( + false, std::memory_order_relaxed); + impl_->batch = + std::make_unique(); + return common::Status::OK(); +} + +common::Status +D3D12AcceleratedExternalDataLoader::PrepareTensor( + const Env& env, + const std::filesystem::path& data_file_path, + std::string_view tensor_name, + FileOffsetType data_offset, + SafeInt data_length) const { + ORT_RETURN_IF_NOT( + impl_->batch != nullptr && + !impl_->batch->finalized, + "D3D12 accelerated initializer batch has not been started."); + return PrepareTensorForBatch( + *impl_->batch, env, data_file_path, + tensor_name, data_offset, data_length); +} + +common::Status +D3D12AcceleratedExternalDataLoader::FinalizeLoad( + const std::function& is_canceled) const { + if (!impl_->enabled) { + return common::Status::OK(); + } + ORT_RETURN_IF_NOT( + impl_->batch != nullptr && + !impl_->batch->finalized, + "D3D12 accelerated initializer batch has not been started."); + auto& batch = *impl_->batch; + const auto fail_or_fallback = + [this, &is_canceled]( + const common::Status& status) + -> common::Status { + if (status.Code() == + common::MODEL_LOAD_CANCELED || + (is_canceled && is_canceled())) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled."); + } + if (IsWeightLoadAccelerationRequired( + impl_->mode)) { + return status; + } + LOGS_DEFAULT(WARNING) + << "D3D12 accelerated initializer loading failed; " + "using the ordinary WebGPU initializer path: " + << status.ErrorMessage(); + impl_->enabled = false; + AbortLoad(); + return common::Status::OK(); + }; + + if (batch.load_range_count == 0) { + if (is_canceled && is_canceled()) { + return fail_or_fallback(ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled.")); + } + batch.finalized = true; + return common::Status::OK(); + } + + D3D12AcceleratedLoadMetrics load_metrics; + const auto loader_status = impl_->EnsureFileLoader(); + if (!loader_status.IsOK()) { + return fail_or_fallback(loader_status); + } + const auto load_status = LoadBatchToD3D12( + batch, *impl_->file_loader, is_canceled, + impl_->abort_requested, load_metrics); + if (!load_status.IsOK()) { + return fail_or_fallback(load_status); + } + + if (is_canceled && is_canceled()) { + return fail_or_fallback(ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled.")); + } + + const auto import_start = Clock::now(); + const auto import_status = + [&]() -> common::Status { + ORT_RETURN_IF_NOT( + impl_->context.D3D12SharedResourceFeaturesAvailable(), + "D3D12 accelerated external weights require Dawn " + "SharedBufferMemoryD3D12Resource and SharedFenceDXGISharedHandle features."); + for (auto& tensor : batch.tensors) { + if (is_canceled && is_canceled()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled."); + } + if (tensor.key.length == 0) { + continue; + } + + SharedBufferMemoryD3D12ResourceDescriptor + resource_descriptor; + resource_descriptor.resource = tensor.resource; + wgpu::SharedBufferMemoryDescriptor + memory_descriptor{}; + memory_descriptor.label = + tensor.key.name.c_str(); + memory_descriptor.nextInChain = + &resource_descriptor; + tensor.memory = + impl_->context.Device() + .ImportSharedBufferMemory( + &memory_descriptor); + ORT_RETURN_IF_NOT( + tensor.memory, + "Failed to import D3D12 accelerated initializer \"", + tensor.key.name, "\" into Dawn."); + + wgpu::SharedBufferMemoryProperties properties{}; + ORT_RETURN_IF_NOT( + tensor.memory.GetProperties(&properties) == + wgpu::Status::Success, + "Failed to query imported D3D12 accelerated initializer \"", + tensor.key.name, "\"."); + ORT_RETURN_IF( + properties.size < + static_cast( + tensor.key.length), + "Imported D3D12 accelerated initializer \"", + tensor.key.name, + "\" is smaller than its tensor payload."); + + wgpu::BufferDescriptor buffer_descriptor{}; + buffer_descriptor.label = + tensor.key.name.c_str(); + buffer_descriptor.size = + static_cast( + tensor.key.length); + buffer_descriptor.usage = + wgpu::BufferUsage::Storage | + wgpu::BufferUsage::CopySrc | + wgpu::BufferUsage::CopyDst; + tensor.buffer = tensor.memory.CreateBuffer( + &buffer_descriptor); + ORT_RETURN_IF_NOT( + tensor.buffer, + "Failed to create buffer for D3D12 accelerated initializer \"", + tensor.key.name, "\"."); + + wgpu::SharedBufferMemoryBeginAccessDescriptor + access_descriptor{}; + access_descriptor.initialized = true; + ORT_RETURN_IF_NOT( + tensor.memory.BeginAccess( + tensor.buffer, + &access_descriptor) == + wgpu::Status::Success, + "Failed to begin Dawn access for D3D12 accelerated initializer \"", + tensor.key.name, "\"."); + tensor.access_started = true; + } + return common::Status::OK(); + }(); + if (!import_status.IsOK()) { + return fail_or_fallback(import_status); + } + const auto import_end = Clock::now(); + batch.finalized = true; + + LOGS_DEFAULT(INFO) + << "WebGPU D3D12 accelerated external initializer load: " + << "file-to-D3D12=" << load_metrics.load_ms + << " ms, Dawn import/access=" + << Milliseconds(import_end - import_start) + << " ms, bytes=" << batch.total_bytes + << ", tensors=" << batch.tensors.size() + << ", loaded_ranges=" + << load_metrics.loaded_ranges << "."; + return common::Status::OK(); +} + +void D3D12AcceleratedExternalDataLoader::AbortLoad() + const noexcept { + if (impl_) { + impl_->abort_requested.store( + true, std::memory_order_relaxed); + } + if (!impl_) { + return; + } + + ORT_TRY { + if (impl_->batch) { + for (auto& tensor : + impl_->batch->tensors) { + if (!tensor.claimed && + tensor.access_started) { + ImportedAllocation allocation; + allocation.resource = + std::move(tensor.resource); + allocation.heap_owner = + std::move(tensor.heap_owner); + allocation.memory = + std::move(tensor.memory); + allocation.buffer = + std::move(tensor.buffer); + allocation.access_started = true; + EndAccessNoThrow(allocation); + tensor.access_started = false; + } + } + impl_->batch.reset(); + } + } + ORT_CATCH(...) { + // Abort is best-effort and is required not to throw. + } +} + +common::Status +D3D12AcceleratedExternalDataLoader::LoadTensor( + const Env&, + const std::filesystem::path& data_file_path, + std::string_view tensor_name, + FileOffsetType data_offset, + SafeInt data_length, + const std::shared_ptr& allocator, + Tensor& tensor) const { + ORT_RETURN_IF_NOT( + impl_->batch != nullptr && + impl_->batch->finalized, + "D3D12 accelerated initializer batch has not been finalized."); + ORT_RETURN_IF( + data_offset < 0, + "D3D12 accelerated initializer has a negative file offset."); + ORT_RETURN_IF_NOT( + allocator != nullptr, + "D3D12 accelerated initializer requires its device allocator."); + ORT_RETURN_IF_NOT( + allocator.get() == + impl_->state->impl_->allocator, + "D3D12 accelerated initializer was passed a different allocator."); + + const size_t length = + static_cast(data_length); + const auto file_iterator = + impl_->batch->files.find( + data_file_path.lexically_normal()); + ORT_RETURN_IF( + file_iterator == impl_->batch->files.end(), + "No prepared D3D12 data file matches \"", + data_file_path.string(), "\"."); + const TensorKey key = MakeKey( + file_iterator->second.canonical_path, + tensor_name, + static_cast(data_offset), + length); + const auto iterator = + impl_->batch->tensors_by_key.find(key); + ORT_RETURN_IF( + iterator == + impl_->batch->tensors_by_key.end(), + "No prepared D3D12 accelerated initializer matches \"", + tensor_name, + "\" and the requested file range."); + + auto& prepared = + impl_->batch->tensors[iterator->second]; + ORT_RETURN_IF( + prepared.claimed, + "D3D12 accelerated initializer \"", + tensor_name, + "\" has already been consumed."); + ORT_RETURN_IF( + length != tensor.SizeInBytes(), + "D3D12 accelerated initializer \"", + tensor_name, + "\" length does not match the placeholder tensor."); + if (length == 0) { + prepared.claimed = true; + tensor = Tensor{ + tensor.DataType(), tensor.Shape(), + nullptr, allocator}; + return common::Status::OK(); + } + ORT_RETURN_IF_NOT( + prepared.buffer && + prepared.access_started, + "D3D12 accelerated initializer \"", + tensor_name, + "\" does not have an imported buffer."); + + auto imported = + std::make_unique(); + imported->resource = + std::move(prepared.resource); + imported->heap_owner = + std::move(prepared.heap_owner); + imported->memory = + std::move(prepared.memory); + imported->buffer = + std::move(prepared.buffer); + imported->access_started = + prepared.access_started; + WGPUBuffer buffer = imported->buffer.Get(); + + { + std::lock_guard lock{ + impl_->state->impl_->mutex}; + ORT_RETURN_IF( + impl_->state->impl_->imported_allocations + .find(buffer) != + impl_->state->impl_->imported_allocations + .end(), + "D3D12 accelerated imported buffer was registered twice."); + impl_->state->impl_->imported_allocations.emplace( + buffer, std::move(imported)); + } + + prepared.access_started = false; + prepared.claimed = true; + tensor = Tensor{ + tensor.DataType(), tensor.Shape(), + buffer, allocator}; + return common::Status::OK(); +} + +} // namespace webgpu +} // namespace onnxruntime + +#endif // defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h new file mode 100644 index 0000000000000..993276cdde9ab --- /dev/null +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + +#include +#include +#include + +#include "core/framework/allocator.h" +#include "core/framework/external_data_loader.h" +#include "core/providers/webgpu/webgpu_provider_options.h" + +struct ID3D12Device; + +namespace onnxruntime { +namespace webgpu { + +class WebGpuContext; +struct CommandRecordingState; + +common::Status CheckD3D12AcceleratedExternalWeightsSupport(WebGpuContext& context); + +common::Status ResolveWeightLoadAccelerationMode( + WeightLoadAccelerationMode mode, + const common::Status& support_status, + bool& enabled); + +// Shared by the loader and allocator so imported resources outlive initializer loading. +class D3D12AcceleratedInitializerState { + public: + ~D3D12AcceleratedInitializerState(); + + private: + struct Impl; + + D3D12AcceleratedInitializerState(); + + std::unique_ptr impl_; + + friend AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + std::function recording_getter, + std::shared_ptr& out_state); + friend class D3D12AcceleratedExternalDataLoader; + friend class D3D12AcceleratedWebGpuAllocator; +}; + +AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + std::function recording_getter, + std::shared_ptr& out_state); + +class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { + public: + D3D12AcceleratedExternalDataLoader( + WebGpuContext& context, + std::shared_ptr state, + WeightLoadAccelerationMode mode); + ~D3D12AcceleratedExternalDataLoader() override; + + bool CanLoad(const OrtMemoryInfo& target_memory_info) const override; + bool SupportsDataType(int32_t tensor_data_type) const override; + bool CreatesTensorForDevice(const OrtDevice& target_device) const override; + common::Status BeginLoad() const override; + common::Status PrepareTensor(const Env& env, + const std::filesystem::path& data_file_path, + std::string_view tensor_name, + FileOffsetType data_offset, + SafeInt data_length) const override; + common::Status FinalizeLoad(const std::function& is_canceled) const override; + void AbortLoad() const noexcept override; + common::Status LoadTensor(const Env& env, + const std::filesystem::path& data_file_path, + std::string_view tensor_name, + FileOffsetType data_offset, + SafeInt data_length, + const std::shared_ptr& allocator, + Tensor& tensor) const override; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace webgpu +} // namespace onnxruntime + +#endif // defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 4411c922a6fc3..a4681c788156a 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -16,26 +16,27 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data, size_t bytes) const { auto& command_state = recording_; if (bytes > 0) { + const auto& buffer_manager = buffer_manager_getter_(); if (dst_is_gpu) { if (src_is_gpu) { // copy from GPU to GPU - buffer_manager_.MemCpy(command_state, - static_cast(const_cast(src_data)), - static_cast(dst_data), - bytes); + buffer_manager.MemCpy(command_state, + static_cast(const_cast(src_data)), + static_cast(dst_data), + bytes); } else { // copy from CPU to GPU - buffer_manager_.Upload(command_state, - const_cast(src_data), - static_cast(dst_data), - bytes); + buffer_manager.Upload(command_state, + const_cast(src_data), + static_cast(dst_data), + bytes); } } else { // copy from GPU to CPU - buffer_manager_.Download(command_state, - static_cast(const_cast(src_data)), - dst_data, - bytes); + buffer_manager.Download(command_state, + static_cast(const_cast(src_data)), + dst_data, + bytes); } } diff --git a/onnxruntime/core/providers/webgpu/data_transfer.h b/onnxruntime/core/providers/webgpu/data_transfer.h index 1ebc869f66d39..b5cc91d2a4b17 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.h +++ b/onnxruntime/core/providers/webgpu/data_transfer.h @@ -3,6 +3,8 @@ #pragma once +#include + #include "core/common/status.h" #include "core/framework/data_transfer.h" @@ -56,8 +58,8 @@ struct CommandRecordingState; // DataTransfer (IDataTransfer subclass) and the C API wrapper. class DataTransferImpl { public: - DataTransferImpl(const BufferManager& buffer_manager, CommandRecordingState& recording) - : buffer_manager_{buffer_manager}, recording_{recording} {} + DataTransferImpl(std::function buffer_manager_getter, CommandRecordingState& recording) + : buffer_manager_getter_{std::move(buffer_manager_getter)}, recording_{recording} {} common::Status CopyTensor(void const* src_data, bool src_is_gpu, @@ -66,14 +68,16 @@ class DataTransferImpl { size_t bytes) const; private: - const BufferManager& buffer_manager_; + std::function buffer_manager_getter_; CommandRecordingState& recording_; }; class DataTransfer : public IDataTransfer { public: - DataTransfer(const BufferManager& buffer_manager, CommandRecordingState& recording, int context_id) - : impl_{buffer_manager, recording}, context_id_{context_id} {} + DataTransfer(std::function buffer_manager_getter, + CommandRecordingState& recording, + int context_id) + : impl_{std::move(buffer_manager_getter), recording}, context_id_{context_id} {} ~DataTransfer() {}; // Device-compatibility half of CanCopy, split out because it needs no BufferManager and so can diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc index f7831672483b3..ac60452990523 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc +++ b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc @@ -83,7 +83,7 @@ CommandRecordingState& GetWebGpuStreamCommandState(const OrtSyncStream* stream) common::Status CopyTensorOnWebGpuStream(const OrtSyncStream* stream, const void* src_data, bool src_is_gpu, void* dst_data, bool dst_is_gpu, size_t bytes) { auto& ep = WebGpuSyncStream::From(stream).ep_; - DataTransferImpl transfer(ep.BufferManager(), ep.Recording()); + DataTransferImpl transfer([&ep]() -> const BufferManager& { return ep.BufferManager(); }, ep.Recording()); return transfer.CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, bytes); } diff --git a/onnxruntime/core/providers/webgpu/external_data_loader.cc b/onnxruntime/core/providers/webgpu/external_data_loader.cc index 6da9598b146f5..51bc230805b54 100644 --- a/onnxruntime/core/providers/webgpu/external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/external_data_loader.cc @@ -18,8 +18,10 @@ bool ExternalDataLoader::CanLoad(const OrtMemoryInfo& target_memory_info) const common::Status ExternalDataLoader::LoadTensor(const Env& env, const std::filesystem::path& data_file_path, + std::string_view /*tensor_name*/, FileOffsetType data_offset, SafeInt data_length, + const std::shared_ptr& /*allocator*/, Tensor& tensor) const { ExternalDataLoadType load_type; if (tensor.Location().device.Type() == OrtDevice::CPU) { diff --git a/onnxruntime/core/providers/webgpu/external_data_loader.h b/onnxruntime/core/providers/webgpu/external_data_loader.h index 7ced4e930bf7a..15a25bca3dd32 100644 --- a/onnxruntime/core/providers/webgpu/external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/external_data_loader.h @@ -19,8 +19,10 @@ class ExternalDataLoader : public IExternalDataLoader { common::Status LoadTensor(const Env& env, const std::filesystem::path& data_file_path, + std::string_view tensor_name, FileOffsetType data_offset, SafeInt data_length, + const std::shared_ptr& allocator, Tensor& tensor) const override; }; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index cc5567b9ce6bf..eafacdf6b4d47 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -3,11 +3,20 @@ #include #include +#include #include #include #include #include +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +#include +#include +#include + +#include "core/providers/webgpu/d3d12_external_data_loader.h" +#endif + #if defined(__GNUC__) #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wstrict-aliasing" @@ -42,256 +51,26 @@ namespace onnxruntime { namespace webgpu { -void WebGpuContext::Initialize(const WebGpuContextConfig& config) { - std::call_once(init_flag_, [this, &config]() { - max_num_pending_dispatches_ = config.max_num_pending_dispatches; - enable_robustness_ = config.enable_robustness; - adapter_index_ = config.adapter_index; - adapter_power_preference_ = config.power_preference; - adapter_backend_type_ = config.backend_type; - - ORT_ENFORCE(!config.adapter_index || config.device == nullptr, - "adapterIndex cannot be used with an externally supplied WebGPU device."); - -#if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) - ORT_ENFORCE(!config.adapter_index, - "adapterIndex requires a native Dawn build with adapter enumeration support."); -#endif - - // Three easily-conflated concepts, at three layers (a pipeline, not the same flag): - // * allow_virtual_devices (env) -- selectability: surface a virtual GPU OrtEpDevice so WebGPU is - // pickable when OS enumeration finds no GPU (e.g. Win32k sandbox). - // * compile_only (session) -- intent: transform only, never finalize/run. - // * device-free / HasDevice() (ctx) -- mechanism: no Dawn device, no-op allocator. - // compile_only alone is valid (device-free even with a real GPU); a virtual device without compile_only is - // rejected at factory CreateEp -- it would try to build a real Dawn device with no hardware. - if (config.compile_only) { - // Device-free: skip Dawn adapter/device creation. Such a context only transforms the graph; the session - // stops before finalization and never executes kernels or allocates. - LOGS_DEFAULT(INFO) << "WebGPU EP context created device-free (compile-only session, no Dawn device)."; - return; - } - - if (device_ == nullptr) { - // Create wgpu::Adapter - wgpu::RequestAdapterOptions req_adapter_options = {}; - req_adapter_options.backendType = static_cast(config.backend_type); - req_adapter_options.powerPreference = static_cast(config.power_preference); - -#if !defined(__wasm__) - auto enabled_adapter_toggles = GetEnabledAdapterToggles(); - - wgpu::DawnTogglesDescriptor adapter_toggles_desc = {}; - adapter_toggles_desc.enabledToggleCount = enabled_adapter_toggles.size(); - adapter_toggles_desc.enabledToggles = enabled_adapter_toggles.data(); - - req_adapter_options.nextInChain = &adapter_toggles_desc; -#endif - - wgpu::Adapter adapter; - if (config.adapter_index) { -#if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) - const auto adapters = EnumerateBundledDawnAdapters(instance_.Get(), req_adapter_options); - ORT_ENFORCE(*config.adapter_index < adapters.size(), - "WebGPU adapterIndex ", *config.adapter_index, - " is out of range; Dawn enumerated ", adapters.size(), - " adapter(s) for the requested backend and power-preference hint."); - adapter = adapters[*config.adapter_index]; - LOGS_DEFAULT(INFO) << "WebGPU EP selected physical adapter index " << *config.adapter_index - << " of " << adapters.size() - << " adapter(s) for the requested backend and power-preference hint."; -#endif - } else { - // Capture adapter request result without throwing inside the Dawn callback. - // Throwing C++ exceptions inside Dawn callbacks leaves Dawn's internal mutexes locked, - // which causes a self-deadlock when the WGPUInstance is later released (e.g., during - // OrtEnv teardown via EventManager::ShutDown()). - struct RequestAdapterResult { - wgpu::RequestAdapterStatus status = wgpu::RequestAdapterStatus::Error; - wgpu::Adapter adapter; - std::string message; - }; - RequestAdapterResult adapter_result; - ORT_ENFORCE(wgpu::WaitStatus::Success == instance_.WaitAny(instance_.RequestAdapter( - &req_adapter_options, - wgpu::CallbackMode::WaitAnyOnly, - [](wgpu::RequestAdapterStatus status, wgpu::Adapter requested_adapter, wgpu::StringView message, - RequestAdapterResult* result) noexcept { - result->status = status; - if (status == wgpu::RequestAdapterStatus::Success) { - result->adapter = std::move(requested_adapter); - } else { - result->message = std::string{message}; - } - }, - &adapter_result), - UINT64_MAX)); - ORT_ENFORCE(adapter_result.status == wgpu::RequestAdapterStatus::Success, - "Failed to get a WebGPU adapter: ", adapter_result.message); - adapter = std::move(adapter_result.adapter); - } - ORT_ENFORCE(adapter != nullptr, "Failed to get a WebGPU adapter."); - - // Create wgpu::Device - wgpu::DeviceDescriptor device_desc = {}; +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +namespace { -#if !defined(__wasm__) - wgpu::DawnTogglesDescriptor device_toggles_desc = {}; - device_desc.nextInChain = &device_toggles_desc; +struct RequestAdapterOptionsLuid : wgpu::ChainedStruct { + RequestAdapterOptionsLuid() { + sType = static_cast(WGPUSType_RequestAdapterOptionsLUID); + } - auto enabled_device_toggles = GetEnabledDeviceToggles(); - device_toggles_desc.enabledToggleCount = enabled_device_toggles.size(); - device_toggles_desc.enabledToggles = enabled_device_toggles.data(); + LUID adapter_luid{}; +}; - auto disabled_device_toggles = GetDisabledDeviceToggles(); - device_toggles_desc.disabledToggleCount = disabled_device_toggles.size(); - device_toggles_desc.disabledToggles = disabled_device_toggles.data(); +} // namespace #endif - std::vector required_features = GetAvailableRequiredFeatures(adapter); - if (!required_features.empty()) { - device_desc.requiredFeatures = required_features.data(); - device_desc.requiredFeatureCount = required_features.size(); - } - wgpu::Limits required_limits = GetRequiredLimits(adapter); - device_desc.requiredLimits = &required_limits; - - device_desc.SetUncapturedErrorCallback( - // Note: Don't throw from a Dawn callback. - [](const wgpu::Device& /*device*/, wgpu::ErrorType type, - wgpu::StringView message, DeviceErrorState* state) noexcept { - { - std::lock_guard lock{state->mutex}; - if (state->status.IsOK()) { - state->status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU device error(", int(type), - "): ", std::string_view{message}); - } - } - if (logging::LoggingManager::HasDefaultLogger()) { - LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; - } - }, - device_error_state_.get()); - // Dawn stops uncaptured callbacks before invoking the lost callback, which retains their state. - device_desc.SetDeviceLostCallback( - wgpu::CallbackMode::AllowSpontaneous, - // Note: Don't throw from a Dawn callback. - [state = device_error_state_](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, - wgpu::StringView message) noexcept { - { - std::lock_guard lock{state->mutex}; - if (state->status.IsOK()) { - state->status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU device lost (", int(reason), - "): ", std::string_view{message}); - } - } - if (logging::LoggingManager::HasDefaultLogger()) { - LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; - } - }); - - struct RequestDeviceResult { - wgpu::RequestDeviceStatus status = wgpu::RequestDeviceStatus::Error; - wgpu::Device device; - std::string message; - }; - RequestDeviceResult device_result; - ORT_ENFORCE(wgpu::WaitStatus::Success == instance_.WaitAny(adapter.RequestDevice( - &device_desc, - wgpu::CallbackMode::WaitAnyOnly, - // Note: Don't throw from a Dawn callback. - [](wgpu::RequestDeviceStatus status, - wgpu::Device device, - wgpu::StringView message, - RequestDeviceResult* result) noexcept { - result->status = status; - if (status == wgpu::RequestDeviceStatus::Success) { - result->device = std::move(device); - } else { - result->message = std::string{message}; - } - }, - &device_result), - UINT64_MAX)); - ORT_ENFORCE(device_result.status == wgpu::RequestDeviceStatus::Success, - "Failed to get a WebGPU device: ", device_result.message); - device_ = std::move(device_result.device); - ORT_ENFORCE(device_ != nullptr, "Failed to get a WebGPU device."); - } - - LOGS_DEFAULT(VERBOSE) << "WebGPU EP Context is created for: Instance=" << instance_.Get() << ", Device=" << device_.Get() << "."; +WebGpuContext::~WebGpuContext() = default; - // cache device queue - device_queue_ = device_.GetQueue(); - // cache device limits - ORT_ENFORCE(Device().GetLimits(&device_limits_) == wgpu::Status::Success); - if (max_storage_buffer_binding_size_ != 0) { - device_limits_.maxStorageBufferBindingSize = - std::min(device_limits_.maxStorageBufferBindingSize, max_storage_buffer_binding_size_); - } - if (max_storage_buffers_per_shader_stage_ != 0) { - ORT_ENFORCE(max_storage_buffers_per_shader_stage_ <= device_limits_.maxStorageBuffersPerShaderStage, - "maxStorageBuffersPerShaderStage exceeds the device limit"); - device_limits_.maxStorageBuffersPerShaderStage = max_storage_buffers_per_shader_stage_; - } - // Align maxStorageBufferBindingSize down to minStorageBufferOffsetAlignment so that - // buffer segment offsets are always properly aligned for WebGPU bind group creation. - if (device_limits_.minStorageBufferOffsetAlignment > 0) { - device_limits_.maxStorageBufferBindingSize -= - (device_limits_.maxStorageBufferBindingSize % device_limits_.minStorageBufferOffsetAlignment); - } - // cache device features - wgpu::SupportedFeatures supported_features; - Device().GetFeatures(&supported_features); - for (size_t i = 0; i < supported_features.featureCount; i++) { - device_features_.insert(supported_features.features[i]); - } -#if !defined(__wasm__) - // Dawn native advertises this software feature on all adapters, and ORT requests it when - // creating a device. An externally supplied device must have requested it too: per-session - // encoders do not protect shared device entry points such as buffer creation and Queue::Submit. - ORT_ENFORCE(DeviceHasFeature(wgpu::FeatureName::ImplicitDeviceSynchronization), - config.device != nullptr - ? "WebGPU: an externally supplied native device must enable ImplicitDeviceSynchronization " - "in DeviceDescriptor.requiredFeatures when it is created." - : "WebGPU: the internally created native device is missing the required " - "ImplicitDeviceSynchronization feature."); -#endif - // cache adapter info - if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { - adapter_info_.nextInChain = &subgroup_matrix_configs_; - } - ORT_ENFORCE(Device().GetAdapterInfo(&adapter_info_) == wgpu::Status::Success); - - // create buffer manager - buffer_mgr_ = BufferManagerFactory::Create(*this, - config.buffer_cache_config.storage.mode, - config.buffer_cache_config.uniform.mode, - config.buffer_cache_config.query_resolve.mode, - config.buffer_cache_config.default_entry.mode); - initializer_buffer_mgr_ = BufferManagerFactory::Create(*this, - BufferCacheMode::LazyRelease, - BufferCacheMode::LazyRelease, - BufferCacheMode::Disabled, - BufferCacheMode::Disabled); - - // create program manager - program_mgr_ = std::make_unique(*this); - - // create split-k config - split_k_config_ = std::make_unique(adapter_info_); - - // set query type -#if !defined(__wasm__) - if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalTimestampQueryInsidePasses)) { - query_type_ = TimestampQueryType::InsidePasses; - } else -#endif - if (DeviceHasFeature(wgpu::FeatureName::TimestampQuery)) { - query_type_ = TimestampQueryType::AtPasses; - } else { - query_type_ = TimestampQueryType::None; - } +void WebGpuContext::StartInitialize(const WebGpuContextConfig& config) { + std::call_once(init_flag_, [this, config]() { + device_free_ = config.compile_only; + Initialize(config); }); if (max_num_pending_dispatches_ != config.max_num_pending_dispatches) { @@ -323,7 +102,7 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { LOGS_DEFAULT(WARNING) << "WebGPU enableRobustness cannot affect an externally supplied WebGPU device. " << "The requested value will be ignored."; - } else if (device_ != nullptr && enable_robustness_ != config.enable_robustness) { + } else if (!device_free_ && enable_robustness_ != config.enable_robustness) { LOGS_DEFAULT(WARNING) << "WebGPU context is already initialized with enableRobustness=" << enable_robustness_ << ". Requested value " << config.enable_robustness << " will be ignored."; @@ -331,6 +110,396 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { } } +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +ID3D12Device* WebGpuContext::WeightLoadingD3D12Device() { + return weight_loading_d3d12_device_.Get(); +} +#endif + +void WebGpuContext::Initialize(const WebGpuContextConfig& config) { + max_num_pending_dispatches_ = config.max_num_pending_dispatches; + enable_robustness_ = config.enable_robustness; + adapter_index_ = config.adapter_index; + adapter_power_preference_ = config.power_preference; + adapter_backend_type_ = config.backend_type; + + ORT_ENFORCE(!config.adapter_index || config.device == nullptr, + "adapterIndex cannot be used with an externally supplied WebGPU device."); +#if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) + ORT_ENFORCE(!config.adapter_index, + "adapterIndex requires a native Dawn build with adapter enumeration support."); +#endif + weight_load_acceleration_mode_ = config.weight_load_acceleration_mode; +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + requested_backend_type_ = static_cast(config.backend_type); +#endif + + // Three easily-conflated concepts, at three layers (a pipeline, not the same flag): + // * allow_virtual_devices (env) -- selectability: surface a virtual GPU OrtEpDevice so WebGPU is + // pickable when OS enumeration finds no GPU (e.g. Win32k sandbox). + // * compile_only (session) -- intent: transform only, never finalize/run. + // * device-free / HasDevice() (ctx) -- mechanism: no Dawn device, no-op allocator. + // compile_only alone is valid (device-free even with a real GPU); a virtual device without compile_only is + // rejected at factory CreateEp -- it would try to build a real Dawn device with no hardware. + if (config.compile_only) { + device_free_ = true; + LOGS_DEFAULT(INFO) << "WebGPU EP context created device-free (compile-only session, no Dawn device)."; + return; + } + + if (device_ == nullptr) { + // Create wgpu::Adapter + wgpu::RequestAdapterOptions req_adapter_options = {}; + req_adapter_options.backendType = static_cast(config.backend_type); + req_adapter_options.powerPreference = static_cast(config.power_preference); + +#if !defined(__wasm__) + auto enabled_adapter_toggles = GetEnabledAdapterToggles(); + + wgpu::DawnTogglesDescriptor adapter_toggles_desc = {}; + adapter_toggles_desc.enabledToggleCount = enabled_adapter_toggles.size(); + adapter_toggles_desc.enabledToggles = enabled_adapter_toggles.data(); + + req_adapter_options.nextInChain = &adapter_toggles_desc; +#endif + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + const auto* ordinary_adapter_options_chain = + req_adapter_options.nextInChain; + RequestAdapterOptionsLuid luid_options{}; + const bool preselect_weight_loading_adapter = + IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_) && + requested_backend_type_ == wgpu::BackendType::D3D12; + if (preselect_weight_loading_adapter) { + LUID selected_luid{}; + const auto selection_start = std::chrono::steady_clock::now(); + const auto selection_status = [&]() -> common::Status { + Microsoft::WRL::ComPtr dxgi_factory; + HRESULT result = + CreateDXGIFactory2(0, IID_PPV_ARGS(&dxgi_factory)); + ORT_RETURN_IF_ERROR( + FAILED(result) + ? ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Accelerated weight loading failed to create the DXGI factory " + "with HRESULT 0x", + std::hex, static_cast(result), ".") + : common::Status::OK()); + + DXGI_GPU_PREFERENCE gpu_preference = + DXGI_GPU_PREFERENCE_UNSPECIFIED; + if (req_adapter_options.powerPreference == + wgpu::PowerPreference::HighPerformance) { + gpu_preference = DXGI_GPU_PREFERENCE_HIGH_PERFORMANCE; + } else if (req_adapter_options.powerPreference == + wgpu::PowerPreference::LowPower) { + gpu_preference = DXGI_GPU_PREFERENCE_MINIMUM_POWER; + } + + for (uint32_t adapter_index = 0;; ++adapter_index) { + Microsoft::WRL::ComPtr dxgi_adapter; + result = dxgi_factory->EnumAdapterByGpuPreference( + adapter_index, gpu_preference, + IID_PPV_ARGS(&dxgi_adapter)); + if (result == DXGI_ERROR_NOT_FOUND) { + break; + } + ORT_RETURN_IF( + FAILED(result), + "Accelerated weight loading failed to enumerate DXGI adapters with " + "HRESULT 0x", + std::hex, static_cast(result), "."); + + DXGI_ADAPTER_DESC1 adapter_description{}; + result = dxgi_adapter->GetDesc1(&adapter_description); + if (FAILED(result) || + (adapter_description.Flags & DXGI_ADAPTER_FLAG_SOFTWARE) != 0) { + continue; + } + + Microsoft::WRL::ComPtr candidate_device; + result = D3D12CreateDevice( + dxgi_adapter.Get(), D3D_FEATURE_LEVEL_11_0, + IID_PPV_ARGS(&candidate_device)); + if (SUCCEEDED(result)) { + selected_luid = adapter_description.AdapterLuid; + weight_loading_d3d12_device_ = std::move(candidate_device); + return common::Status::OK(); + } + } + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Accelerated weight loading did not find a hardware D3D12 adapter."); + }(); + + if (!selection_status.IsOK()) { + if (IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_)) { + ORT_THROW(selection_status.ErrorMessage()); + } + LOGS_DEFAULT(WARNING) + << selection_status.ErrorMessage() + << " Continuing with ordinary Dawn adapter selection without " + "accelerated weight loading."; + weight_loading_d3d12_device_.Reset(); + } + if (weight_loading_d3d12_device_) { + luid_options.adapter_luid = selected_luid; +#if !defined(__wasm__) + luid_options.nextInChain = &adapter_toggles_desc; +#endif + req_adapter_options.nextInChain = &luid_options; + + LOGS_DEFAULT(INFO) + << "WebGPU D3D12 weight-loading GPU selection: " + << std::chrono::duration( + std::chrono::steady_clock::now() - selection_start) + .count() + << " ms."; + } + } +#endif + + wgpu::Adapter adapter; + if (config.adapter_index) { +#if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) + const auto adapters = EnumerateBundledDawnAdapters(instance_.Get(), req_adapter_options); + ORT_ENFORCE(*config.adapter_index < adapters.size(), + "WebGPU adapterIndex ", *config.adapter_index, + " is out of range; Dawn enumerated ", adapters.size(), + " adapter(s) for the requested backend and power-preference hint."); + adapter = adapters[*config.adapter_index]; + LOGS_DEFAULT(INFO) << "WebGPU EP selected physical adapter index " << *config.adapter_index + << " of " << adapters.size() + << " adapter(s) for the requested backend and power-preference hint."; +#endif + } else { + // Capture adapter request result without throwing inside the Dawn callback. + // Throwing C++ exceptions inside Dawn callbacks leaves Dawn's internal mutexes locked, + // which causes a self-deadlock when the WGPUInstance is later released (e.g., during + // OrtEnv teardown via EventManager::ShutDown()). + struct RequestAdapterResult { + wgpu::RequestAdapterStatus status = wgpu::RequestAdapterStatus::Error; + wgpu::Adapter adapter; + std::string message; + }; + const auto request_adapter = [&]() { + RequestAdapterResult result; + ORT_ENFORCE(wgpu::WaitStatus::Success == + instance_.WaitAny(instance_.RequestAdapter( + &req_adapter_options, + wgpu::CallbackMode::WaitAnyOnly, + [](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter, + wgpu::StringView message, + RequestAdapterResult* result) noexcept { + result->status = status; + if (status == wgpu::RequestAdapterStatus::Success) { + result->adapter = std::move(adapter); + } else { + result->message = std::string{message}; + } + }, + &result), + UINT64_MAX)); + return result; + }; + RequestAdapterResult adapter_result = request_adapter(); +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + if (adapter_result.status != wgpu::RequestAdapterStatus::Success && + weight_loading_d3d12_device_ && + !IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_)) { + LOGS_DEFAULT(WARNING) + << "Dawn rejected the adapter selected for accelerated weight " + "loading: " + << adapter_result.message + << " Retrying ordinary Dawn adapter selection without " + "accelerated weight loading."; + weight_loading_d3d12_device_.Reset(); + req_adapter_options.nextInChain = ordinary_adapter_options_chain; + adapter_result = request_adapter(); + } +#endif + ORT_ENFORCE(adapter_result.status == wgpu::RequestAdapterStatus::Success, + "Failed to get a WebGPU adapter: ", adapter_result.message); + adapter = std::move(adapter_result.adapter); + } + ORT_ENFORCE(adapter != nullptr, "Failed to get a WebGPU adapter."); + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + d3d12_shared_resource_features_available_ = + adapter.HasFeature(wgpu::FeatureName::SharedBufferMemoryD3D12Resource) && + adapter.HasFeature(wgpu::FeatureName::SharedFenceDXGISharedHandle); +#endif + // Create wgpu::Device + wgpu::DeviceDescriptor device_desc = {}; + +#if !defined(__wasm__) + wgpu::DawnTogglesDescriptor device_toggles_desc = {}; + device_desc.nextInChain = &device_toggles_desc; + + auto enabled_device_toggles = GetEnabledDeviceToggles(); + device_toggles_desc.enabledToggleCount = enabled_device_toggles.size(); + device_toggles_desc.enabledToggles = enabled_device_toggles.data(); + + auto disabled_device_toggles = GetDisabledDeviceToggles(); + device_toggles_desc.disabledToggleCount = disabled_device_toggles.size(); + device_toggles_desc.disabledToggles = disabled_device_toggles.data(); +#endif + + std::vector required_features = GetAvailableRequiredFeatures(adapter); + if (!required_features.empty()) { + device_desc.requiredFeatures = required_features.data(); + device_desc.requiredFeatureCount = required_features.size(); + } + wgpu::Limits required_limits = GetRequiredLimits(adapter); + device_desc.requiredLimits = &required_limits; + + // TODO: revise temporary error handling + device_desc.SetUncapturedErrorCallback( + // Note: Don't throw from a Dawn callback. + [](const wgpu::Device& /*device*/, wgpu::ErrorType type, + wgpu::StringView message) noexcept { + if (logging::LoggingManager::HasDefaultLogger()) { + LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; + } + }); + // TODO: revise temporary device lost handling + device_desc.SetDeviceLostCallback( + wgpu::CallbackMode::AllowSpontaneous, + // Note: Don't throw from a Dawn callback. + [](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, wgpu::StringView message) noexcept { + if (logging::LoggingManager::HasDefaultLogger()) { + LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; + } + }); + + struct RequestDeviceResult { + wgpu::RequestDeviceStatus status = wgpu::RequestDeviceStatus::Error; + wgpu::Device device; + std::string message; + }; + RequestDeviceResult device_result; + ORT_ENFORCE(wgpu::WaitStatus::Success == instance_.WaitAny(adapter.RequestDevice( + &device_desc, + wgpu::CallbackMode::WaitAnyOnly, + // Note: Don't throw from a Dawn callback. + [](wgpu::RequestDeviceStatus status, + wgpu::Device device, + wgpu::StringView message, + RequestDeviceResult* result) noexcept { + result->status = status; + if (status == wgpu::RequestDeviceStatus::Success) { + result->device = std::move(device); + } else { + result->message = std::string{message}; + } + }, + &device_result), + UINT64_MAX)); + ORT_ENFORCE(device_result.status == wgpu::RequestDeviceStatus::Success, + "Failed to get a WebGPU device: ", device_result.message); + device_ = std::move(device_result.device); + ORT_ENFORCE(device_ != nullptr, "Failed to get a WebGPU device."); + } + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + if (config.device != nullptr) { + wgpu::SupportedFeatures supported_features; + device_.GetFeatures(&supported_features); + bool has_shared_buffer_memory = false; + bool has_shared_fence = false; + for (size_t i = 0; i < supported_features.featureCount; ++i) { + has_shared_buffer_memory |= + supported_features.features[i] == + wgpu::FeatureName::SharedBufferMemoryD3D12Resource; + has_shared_fence |= + supported_features.features[i] == + wgpu::FeatureName::SharedFenceDXGISharedHandle; + } + d3d12_shared_resource_features_available_ = + has_shared_buffer_memory && has_shared_fence; + wgpu::AdapterInfo supplied_adapter_info; + if (device_.GetAdapterInfo(&supplied_adapter_info) == wgpu::Status::Success) { + requested_backend_type_ = supplied_adapter_info.backendType; + } + } +#endif + LOGS_DEFAULT(VERBOSE) << "WebGPU EP Context is created for: Instance=" << instance_.Get() << ", Device=" << device_.Get() << "."; + + // cache device queue + device_queue_ = device_.GetQueue(); + // cache device limits + ORT_ENFORCE(Device().GetLimits(&device_limits_) == wgpu::Status::Success); + if (max_storage_buffer_binding_size_ != 0) { + device_limits_.maxStorageBufferBindingSize = + std::min(device_limits_.maxStorageBufferBindingSize, max_storage_buffer_binding_size_); + } + if (max_storage_buffers_per_shader_stage_ != 0) { + ORT_ENFORCE(max_storage_buffers_per_shader_stage_ <= device_limits_.maxStorageBuffersPerShaderStage, + "maxStorageBuffersPerShaderStage exceeds the device limit"); + device_limits_.maxStorageBuffersPerShaderStage = max_storage_buffers_per_shader_stage_; + } + // Align maxStorageBufferBindingSize down to minStorageBufferOffsetAlignment so that + // buffer segment offsets are always properly aligned for WebGPU bind group creation. + if (device_limits_.minStorageBufferOffsetAlignment > 0) { + device_limits_.maxStorageBufferBindingSize -= + (device_limits_.maxStorageBufferBindingSize % device_limits_.minStorageBufferOffsetAlignment); + } + // cache device features + wgpu::SupportedFeatures supported_features; + Device().GetFeatures(&supported_features); + for (size_t i = 0; i < supported_features.featureCount; i++) { + device_features_.insert(supported_features.features[i]); + } + // cache adapter info +#if !defined(__wasm__) + // Native sessions share device entry points outside their per-session command recordings. + ORT_ENFORCE(DeviceHasFeature(wgpu::FeatureName::ImplicitDeviceSynchronization), + config.device != nullptr + ? "WebGPU: an externally supplied native device must enable ImplicitDeviceSynchronization " + "in DeviceDescriptor.requiredFeatures when it is created." + : "WebGPU: the internally created native device is missing the required " + "ImplicitDeviceSynchronization feature."); + if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { + adapter_info_.nextInChain = &subgroup_matrix_configs_; + } +#endif + ORT_ENFORCE(Device().GetAdapterInfo(&adapter_info_) == wgpu::Status::Success); + + // create buffer manager + buffer_mgr_ = BufferManagerFactory::Create(*this, + config.buffer_cache_config.storage.mode, + config.buffer_cache_config.uniform.mode, + config.buffer_cache_config.query_resolve.mode, + config.buffer_cache_config.default_entry.mode); + + // create initializer buffer manager. + initializer_buffer_mgr_ = BufferManagerFactory::Create(*this, + BufferCacheMode::LazyRelease, + BufferCacheMode::LazyRelease, + BufferCacheMode::Disabled, + BufferCacheMode::Disabled); + + // create program manager + program_mgr_ = std::make_unique(*this); + + // create split-k config + split_k_config_ = std::make_unique(adapter_info_); + + // set query type +#if !defined(__wasm__) + if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalTimestampQueryInsidePasses)) { + query_type_ = TimestampQueryType::InsidePasses; + } else +#endif + if (DeviceHasFeature(wgpu::FeatureName::TimestampQuery)) { + query_type_ = TimestampQueryType::AtPasses; + } else { + query_type_ = TimestampQueryType::None; + } +} + Status WebGpuContext::Wait(wgpu::Future f) { auto status = instance_.WaitAny(f, UINT64_MAX); if (status == wgpu::WaitStatus::Success) { @@ -926,6 +1095,23 @@ std::vector WebGpuContext::GetAvailableRequiredFeatures(const required_features.push_back(feature); } } +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + constexpr wgpu::FeatureName d3d12_import_features[]{ + wgpu::FeatureName::SharedBufferMemoryD3D12Resource, + wgpu::FeatureName::SharedFenceDXGISharedHandle, + }; + for (auto feature : d3d12_import_features) { + if (adapter.HasFeature(feature)) { + required_features.push_back(feature); + } else { + ORT_ENFORCE(!IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_), + "The selected Dawn D3D12 adapter does not support a feature required by " + "weightLoadAcceleration: ", + static_cast(feature)); + } + } +#endif return required_features; } @@ -1255,6 +1441,12 @@ Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, auto command_buffer = recording.command_encoder.Finish(); device_queue_.Submit(1, &command_buffer); refresh_pending_buffers(); +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + for (auto& release : recording.pending_release_callbacks) { + release(); + } + recording.pending_release_callbacks.clear(); +#endif recording.command_encoder = nullptr; recording.num_pending_dispatches = 0; return status; @@ -1441,10 +1633,10 @@ WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& co ORT_ENFORCE(default_instance_ != nullptr, "Failed to create wgpu::Instance."); } - if (context_id == 0) { - // context ID is preserved for the default context. User cannot use context ID 0 as a custom context. + if (context_id == 0 || context_id == kDeviceFreeDefaultContextId) { + // The runnable and device-free default contexts share the default instance but use separate cache entries. ORT_ENFORCE(instance == nullptr && device == nullptr, - "WebGPU EP default context (contextId=0) must not have custom WebGPU instance or device."); + "WebGPU EP default contexts must not have custom WebGPU instance or device."); instance = default_instance_; } else { @@ -1481,7 +1673,7 @@ WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& co // if this was the first (and only) reference, so we don't leave a zombie context in the map // that would later deadlock during Cleanup(). ORT_TRY { - it->second.context->Initialize(config); + it->second.context->StartInitialize(config); } ORT_CATCH(...) { if (--it->second.ref_count == 0) { diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 2fd9e77871e24..c3128ce73927b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include #include @@ -16,6 +17,7 @@ #include "core/common/common.h" #include "core/providers/webgpu/buffer_manager.h" +#include "core/providers/webgpu/webgpu_provider_options.h" #include "core/providers/webgpu/program_manager.h" #include "core/providers/webgpu/webgpu_utils.h" @@ -23,15 +25,21 @@ #include "core/providers/webgpu/webgpu_pix_frame_generator.h" #endif // ENABLE_PIX_FOR_WEBGPU_EP +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +#include +#include +#endif + namespace onnxruntime { class Tensor; namespace webgpu { + +inline constexpr int kDeviceFreeDefaultContextId = -1; class WebGpuContext; class ComputeContextBase; class ComputeContext; class ProgramBase; - // PendingKernelInfo stores profiling information for a kernel execution struct PendingKernelInfo { PendingKernelInfo(std::string_view kernel_name, @@ -122,6 +130,9 @@ struct CommandRecordingState { // Concurrent allocator frees inspect only this flag, never the encoders or deferred dispatches. // BufferManager clears it under its cache lock after submission or abandonment. std::atomic has_unsubmitted_work{false}; +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + std::vector> pending_release_callbacks; +#endif std::vector deferred_dispatches; std::vector pending_kernels; GraphCaptureState graph_capture_state{GraphCaptureState::Default}; @@ -164,6 +175,8 @@ struct WebGpuContextConfig { }; bool enable_robustness_explicitly_set{false}; bool preserve_device{false}; + WeightLoadAccelerationMode weight_load_acceleration_mode{ + WeightLoadAccelerationMode::Off}; // When true, skip Dawn adapter/device creation and all device-dependent initialization; the context // can only be used for graph transformation, not execution. Derived from kOrtSessionOptionCompileOnly. bool compile_only{false}; @@ -251,6 +264,8 @@ class WebGpuContext final { static constexpr uint64_t kWebGpuGuaranteedMaxStorageBufferBindingSize = 128ULL * 1024 * 1024; + ~WebGpuContext(); + Status Wait(wgpu::Future f); const wgpu::Instance& Instance() const { return instance_; } @@ -258,7 +273,9 @@ class WebGpuContext final { const wgpu::AdapterInfo& AdapterInfo() const { return adapter_info_; } const wgpu::Limits& DeviceLimits() const { return device_limits_; } - bool DeviceHasFeature(wgpu::FeatureName feature) const { return device_features_.contains(feature); } + bool DeviceHasFeature(wgpu::FeatureName feature) const { + return device_features_.contains(feature); + } const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return subgroup_matrix_configs_; } const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) { @@ -313,6 +330,7 @@ class WebGpuContext final { // Context-level managers are shared by sessions and synchronize their buffer caches internally. webgpu::BufferManager& BufferManager() const { return *buffer_mgr_; } + webgpu::BufferManager& InitializerBufferManager() const { return *initializer_buffer_mgr_; } inline webgpu::ValidationMode ValidationMode() const { @@ -322,6 +340,16 @@ class WebGpuContext final { // False for a device-free ("virtual device") context, which has no Dawn device and can only run graph // transformation. Used to hand out a no-op allocator instead of a real GpuBufferAllocator. inline bool HasDevice() const { return device_ != nullptr; } + inline bool IsDeviceFree() const { return device_free_; } +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + ID3D12Device* WeightLoadingD3D12Device(); + bool D3D12SharedResourceFeaturesAvailable() const { + return d3d12_shared_resource_features_available_; + } + wgpu::BackendType RequestedBackendType() const { + return requested_backend_type_; + } +#endif // // Get Split-K configuration. @@ -405,6 +433,7 @@ class WebGpuContext final { max_storage_buffers_per_shader_stage_{max_storage_buffers_per_shader_stage} {} ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WebGpuContext); + void StartInitialize(const WebGpuContextConfig& config); void Initialize(const WebGpuContextConfig& config); wgpu::BindGroup CreateBindGroup(const std::vector& bind_buffers, @@ -443,6 +472,7 @@ class WebGpuContext final { friend class WebGpuContextFactory; std::once_flag init_flag_; + bool device_free_ = false; struct DeviceErrorState { std::mutex mutex; @@ -492,6 +522,13 @@ class WebGpuContext final { // Shared GPU profiling events for run-level profiling. profiling::Events events_; bool preserve_device_; + WeightLoadAccelerationMode weight_load_acceleration_mode_{ + WeightLoadAccelerationMode::Off}; +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + Microsoft::WRL::ComPtr weight_loading_d3d12_device_; + bool d3d12_shared_resource_features_available_ = false; + wgpu::BackendType requested_backend_type_ = wgpu::BackendType::Undefined; +#endif uint64_t max_storage_buffer_binding_size_; uint32_t max_storage_buffers_per_shader_stage_; GraphCaptureState graph_capture_state_{GraphCaptureState::Default}; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 690ea8114a200..fd1a267aa1fa8 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -2,6 +2,9 @@ // Licensed under the MIT License. #include "core/providers/webgpu/webgpu_execution_provider.h" +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +#include "core/providers/webgpu/d3d12_external_data_loader.h" +#endif #include #include @@ -71,7 +74,9 @@ class Memcpy final : public OpKernel { const auto* X = ctx->Input(0); Tensor* Y = ctx->Output(0, X->Shape()); const auto& ep = *static_cast(Info().GetExecutionProvider()); - DataTransfer transfer(ep.BufferManager(), ep.Recording(), ep.GetDeviceId()); + DataTransfer transfer([&ep]() -> const BufferManager& { return ep.BufferManager(); }, + ep.Recording(), + ep.GetDeviceId()); return transfer.CopyTensor(*X, *Y); } }; @@ -617,12 +622,38 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, multi_rotary_cache_concat_offset_{config.multi_rotary_cache_concat_offset}, kv_cache_quantization_bits_{config.kv_cache_quantization_bits}, enable_matmul_fp32_accumulation_{config.enable_matmul_fp32_accumulation}, +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + weight_load_acceleration_mode_{config.weight_load_acceleration_mode}, +#endif recording_{std::make_unique()}, prepack_allocator_{CreateWebGpuAllocator( context_id, /*device_free=*/!context.HasDevice(), [this]() -> const webgpu::BufferManager& { return InitializerBufferManager(); }, [this]() -> webgpu::CommandRecordingState& { return Recording(); }, false)} { +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + if (webgpu::IsWeightLoadAccelerationEnabled( + config.weight_load_acceleration_mode)) { + accelerated_initializer_allocator_ = + CreateD3D12AcceleratedWebGpuAllocator( + context_, + [this]() -> webgpu::CommandRecordingState& { return Recording(); }, + accelerated_initializer_state_); + } +#else + if (webgpu::IsWeightLoadAccelerationRequired( + config.weight_load_acceleration_mode)) { + ORT_THROW( + "The requested weightLoadAcceleration mode requires a supported " + "disk-to-GPU weight loading implementation."); + } + if (webgpu::IsWeightLoadAccelerationEnabled( + config.weight_load_acceleration_mode)) { + LOGS_DEFAULT(WARNING) + << "Accelerated weight loading is unavailable in this build; using " + "the ordinary WebGPU initializer loading path."; + } +#endif if (enable_graph_capture_ && config.session_buffer_pool_generations > 0) { session_buffer_pool_ = std::make_unique( config.session_buffer_pool_generations); @@ -640,9 +671,14 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, } std::vector WebGpuExecutionProvider::CreatePreferredAllocators() { - const bool device_free = !context_.HasDevice(); + const bool device_free = context_.IsDeviceFree(); return { // allocator for initializers +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + accelerated_initializer_allocator_ != nullptr + ? accelerated_initializer_allocator_ + : +#endif CreateWebGpuAllocator( context_id_, device_free, @@ -763,13 +799,28 @@ std::vector> WebGpuExecutionProvider::GetCapa #endif // !defined(ORT_USE_EP_API_ADAPTERS) std::unique_ptr WebGpuExecutionProvider::GetDataTransfer() const { - return std::make_unique(BufferManager(), Recording(), context_id_); + return std::make_unique( + [&context = context_]() -> const webgpu::BufferManager& { + return context.BufferManager(); + }, + Recording(), + context_id_); } #if defined(__wasm__) std::unique_ptr WebGpuExecutionProvider::GetExternalDataLoader() const { return std::make_unique(); } +#elif defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +std::unique_ptr WebGpuExecutionProvider::GetExternalDataLoader() const { + if (accelerated_initializer_state_ == nullptr) { + return nullptr; + } + + return std::make_unique( + context_, accelerated_initializer_state_, + weight_load_acceleration_mode_); +} #endif std::optional WebGpuExecutionProvider::ShouldConvertDataLayoutForOp(std::string_view node_domain, diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index 507b426cb1ec1..caeb5a1a2ca60 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -18,6 +18,7 @@ #include "core/providers/providers.h" #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/session_buffer_pool.h" +#include "core/providers/webgpu/webgpu_provider_options.h" #if defined(ENABLE_PIX_FOR_WEBGPU_EP) #include "core/providers/webgpu/webgpu_pix_frame_generator.h" @@ -40,6 +41,9 @@ KernelCreateInfo BuildKernelCreateInfo(); class WebGpuContext; class WebGpuProfiler; class GpuBufferAllocator; +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +class D3D12AcceleratedInitializerState; +#endif // Forward declare CapturedCommandInfo which is now defined in webgpu_context.h struct CapturedCommandInfo; @@ -49,8 +53,10 @@ std::shared_ptr GetKernelRegistry(bool enable_graph_capture, boo } // namespace webgpu struct WebGpuExecutionProviderConfig { - DataLayout data_layout{DataLayout::NHWC}; // preferred layout is NHWC by default - bool enable_graph_capture{false}; // graph capture feature is disabled by default + DataLayout data_layout{DataLayout::NHWC}; // preferred layout is NHWC by default + bool enable_graph_capture{false}; // graph capture feature is disabled by default + webgpu::WeightLoadAccelerationMode weight_load_acceleration_mode{ + webgpu::WeightLoadAccelerationMode::Off}; bool enable_pix_capture{false}; // PIX capture is disabled by default bool enable_int64{false}; // int64 ops are not enabled by default uint32_t multi_rotary_cache_concat_offset{0}; // offset for concatenated multi rotary cache (0 = disabled) @@ -87,7 +93,7 @@ class WebGpuExecutionProvider : public IExecutionProvider { } #endif std::unique_ptr GetDataTransfer() const override; -#if defined(__wasm__) +#if defined(__wasm__) || (defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING)) std::unique_ptr GetExternalDataLoader() const override; #endif @@ -151,6 +157,12 @@ class WebGpuExecutionProvider : public IExecutionProvider { uint32_t multi_rotary_cache_concat_offset_ = 0; uint32_t kv_cache_quantization_bits_ = 0; bool enable_matmul_fp32_accumulation_ = false; +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + webgpu::WeightLoadAccelerationMode weight_load_acceleration_mode_{ + webgpu::WeightLoadAccelerationMode::Off}; + std::shared_ptr accelerated_initializer_state_; + AllocatorPtr accelerated_initializer_allocator_; +#endif std::unordered_map graph_id_to_run_count_; // Required regular runs before graph capture for any necessary allocations. const int min_num_runs_before_graph_capture_ = 0; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 5357323fe46f0..9c79b250f4ae0 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -47,6 +47,24 @@ struct WebGpuProviderFactory : IExecutionProviderFactory { namespace { +WeightLoadAccelerationMode ParseWeightLoadAccelerationMode( + const ConfigOptions& config_options) { + std::string value; + if (!config_options.TryGetConfigEntry(kWeightLoadAcceleration, value) || + value == kWeightLoadAcceleration_Off) { + return WeightLoadAccelerationMode::Off; + } + if (value == kWeightLoadAcceleration_Preferred) { + return WeightLoadAccelerationMode::Preferred; + } + if (value == kWeightLoadAcceleration_Required) { + return WeightLoadAccelerationMode::Required; + } + ORT_THROW( + "Invalid weightLoadAcceleration value: ", value, + ". Must be \"off\", \"preferred\", or \"required\"."); +} + WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) { WebGpuExecutionProviderConfig webgpu_ep_config{}; @@ -72,6 +90,9 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) } } + webgpu_ep_config.weight_load_acceleration_mode = + ParseWeightLoadAccelerationMode(config_options); + if (std::string pool_generations_str; config_options.TryGetConfigEntry(kSessionBufferPoolGenerations, pool_generations_str)) { size_t pool_generations = 0; @@ -167,6 +188,8 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) LOGS_DEFAULT(VERBOSE) << "WebGPU EP preferred layout: " << int(webgpu_ep_config.data_layout); LOGS_DEFAULT(VERBOSE) << "WebGPU EP graph capture enable: " << webgpu_ep_config.enable_graph_capture; + LOGS_DEFAULT(VERBOSE) << "WebGPU EP weight load acceleration mode: " + << static_cast(webgpu_ep_config.weight_load_acceleration_mode); LOGS_DEFAULT(VERBOSE) << "WebGPU EP force CPU node count: " << webgpu_ep_config.force_cpu_node_names.size(); LOGS_DEFAULT(VERBOSE) << "WebGPU EP pix capture enable: " << webgpu_ep_config.enable_pix_capture; LOGS_DEFAULT(VERBOSE) << "WebGPU EP enable int64: " << webgpu_ep_config.enable_int64; @@ -179,14 +202,19 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) WebGpuContextConfig ParseWebGpuContextConfig(const ConfigOptions& config_options) { WebGpuContextConfig config{}; + config.weight_load_acceleration_mode = + ParseWeightLoadAccelerationMode(config_options); + if (std::string context_id_str; config_options.TryGetConfigEntry(kDeviceId, context_id_str)) { const auto result = std::from_chars(context_id_str.data(), context_id_str.data() + context_id_str.size(), - config.context_id); + config.context_id); ORT_ENFORCE(result.ec == std::errc{} && result.ptr == context_id_str.data() + context_id_str.size() && config.context_id >= 0 && config.context_id <= std::numeric_limits::max(), "Invalid deviceId value: ", context_id_str, ". Must be an integer in the range 0 to ", std::numeric_limits::max(), "."); + ORT_ENFORCE(config.context_id != kDeviceFreeDefaultContextId, + "WebGPU device ID ", kDeviceFreeDefaultContextId, " is reserved for internal use."); } if (std::string adapter_index_str; @@ -275,6 +303,9 @@ WebGpuContextConfig ParseWebGpuContextConfig(const ConfigOptions& config_options // is derived from the session config kOrtSessionOptionCompileOnly, which the Compile API sets // automatically -- same signal other EPs use (e.g. NV TensorRT RTX). Not a WebGPU-specific option. config.compile_only = config_options.GetConfigOrDefault(kOrtSessionOptionCompileOnly, "0") == "1"; + if (config.compile_only && config.context_id == 0) { + config.context_id = kDeviceFreeDefaultContextId; + } std::string max_storage_buffer_binding_size_str; if (config_options.TryGetConfigEntry(kMaxStorageBufferBindingSize, max_storage_buffer_binding_size_str)) { diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h index 1da47ace5d96a..6d0569d3688dd 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h @@ -5,12 +5,28 @@ namespace onnxruntime { namespace webgpu { + +enum class WeightLoadAccelerationMode { + Off, + Preferred, + Required, +}; + +constexpr bool IsWeightLoadAccelerationEnabled(WeightLoadAccelerationMode mode) { + return mode != WeightLoadAccelerationMode::Off; +} + +constexpr bool IsWeightLoadAccelerationRequired(WeightLoadAccelerationMode mode) { + return mode == WeightLoadAccelerationMode::Required; +} + namespace options { // The following are the options that can be set in the WebGPU provider options. constexpr const char* kPreferredLayout = "ep.webgpuexecutionprovider.preferredLayout"; constexpr const char* kEnableGraphCapture = "ep.webgpuexecutionprovider.enableGraphCapture"; +constexpr const char* kWeightLoadAcceleration = "ep.webgpuexecutionprovider.weightLoadAcceleration"; // Number of generations of buffers to retain in the per-session pool for reuse // across captured-graph lifetimes. 0 disables pooling. Default 1 caches one // generator's worth of intermediate buffers. @@ -77,6 +93,10 @@ constexpr const char* kPreferredLayout_NHWC = "NHWC"; constexpr const char* kEnableGraphCapture_ON = "1"; constexpr const char* kEnableGraphCapture_OFF = "0"; +constexpr const char* kWeightLoadAcceleration_Off = "off"; +constexpr const char* kWeightLoadAcceleration_Preferred = "preferred"; +constexpr const char* kWeightLoadAcceleration_Required = "required"; + constexpr const char* kEnableInt64_ON = "1"; constexpr const char* kEnableInt64_OFF = "0"; diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 9d80a191ae0da..4ce943a40bb20 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1212,6 +1212,7 @@ common::Status InferenceSession::LoadWithLoader(std::function +#include +#include + +#include "gtest/gtest.h" +#include "test/util/include/test_environment.h" + #if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) #include #include -#include #include #include "core/common/inlined_containers.h" @@ -14,16 +22,17 @@ #include "core/graph/onnx_protobuf.h" #include "core/providers/cpu/cpu_execution_provider.h" #include "core/session/inference_session.h" -#include "gtest/gtest.h" #include "test/test_environment.h" #include "test/unittest_util/framework_test_utils.h" #include "test/util/include/asserts.h" #include "test/util/include/file_util.h" +#endif namespace onnxruntime { namespace test { namespace { +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) enum class ReadFailure { None, Status, Exception }; @@ -297,8 +306,102 @@ TEST_F(ExternalDataLoaderLifetimeTest, RecreatesLoaderAfterFactoryFailure) { } #endif +#endif + +class BatchLifecycleExternalDataLoader final : public IExternalDataLoader { + public: + enum class FailurePoint { + None, + Begin, + Finalize, + }; + + explicit BatchLifecycleExternalDataLoader(FailurePoint failure_point = FailurePoint::None) + : failure_point_{failure_point} { + } + + bool CanLoad(const OrtMemoryInfo&) const override { + return false; + } + + common::Status BeginLoad() const override { + ++begin_count; + return failure_point_ == FailurePoint::Begin + ? ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "begin failure") + : common::Status::OK(); + } + + common::Status FinalizeLoad(const std::function&) const override { + ++finalize_count; + return failure_point_ == FailurePoint::Finalize + ? ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "finalize failure") + : common::Status::OK(); + } + + void AbortLoad() const noexcept override { + ++abort_count; + } + + mutable int begin_count = 0; + mutable int finalize_count = 0; + mutable int abort_count = 0; + + private: + FailurePoint failure_point_; +}; + +TEST(ExternalDataLoaderManagerTest, DefaultBatchLifecycleIsBackwardCompatible) { + class LegacyExternalDataLoader final : public IExternalDataLoader { + public: + bool CanLoad(const OrtMemoryInfo&) const override { + return false; + } + }; + + ExternalDataLoaderManager manager; + ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::make_unique())); + EXPECT_STATUS_OK(manager.BeginLoad()); + EXPECT_STATUS_OK(manager.FinalizeLoad([]() { return false; })); + manager.AbortLoad(); +} + +TEST(ExternalDataLoaderManagerTest, BeginFailureAbortsEveryLoader) { + ExternalDataLoaderManager manager; + auto first = std::make_unique(); + auto* first_ptr = first.get(); + auto failing = std::make_unique( + BatchLifecycleExternalDataLoader::FailurePoint::Begin); + auto* failing_ptr = failing.get(); + + ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(first))); + ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(failing))); + + EXPECT_FALSE(manager.BeginLoad().IsOK()); + EXPECT_EQ(first_ptr->begin_count, 1); + EXPECT_EQ(failing_ptr->begin_count, 1); + EXPECT_EQ(first_ptr->abort_count, 1); + EXPECT_EQ(failing_ptr->abort_count, 1); +} + +TEST(ExternalDataLoaderManagerTest, FinalizeFailureAbortsEveryLoader) { + ExternalDataLoaderManager manager; + auto failing = std::make_unique( + BatchLifecycleExternalDataLoader::FailurePoint::Finalize); + auto* failing_ptr = failing.get(); + auto unfinalized = std::make_unique(); + auto* unfinalized_ptr = unfinalized.get(); + + ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(failing))); + ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(unfinalized))); + + ASSERT_STATUS_OK(manager.BeginLoad()); + EXPECT_FALSE(manager.FinalizeLoad([]() { return false; }).IsOK()); + EXPECT_EQ(failing_ptr->finalize_count, 1); + EXPECT_EQ(unfinalized_ptr->finalize_count, 0); + EXPECT_EQ(failing_ptr->abort_count, 1); + EXPECT_EQ(unfinalized_ptr->abort_count, 1); +} + } // namespace } // namespace test } // namespace onnxruntime - -#endif diff --git a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc index 21363f82f2752..8a6c9b5186a37 100644 --- a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc +++ b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc @@ -76,6 +76,454 @@ TEST(WebGpuConcurrentContextTestStandalone, ShaderDumpWritesRemainComplete) { } } +// Fixture: builds the shared model once and provides session/feed factories. A long-lived +// keepalive session pins the shared WebGPU context (ref-count > 0) for the whole test so that +// churn/destroy in one thread never tears the context down under the others. +class WebGpuConcurrentContextTest : public ::testing::Test { + protected: + static constexpr int64_t kNumElements = 256 * 1024; // 1 MB per initializer + static constexpr int kChainLen = 8; // 8 uploads + 8 dispatches per session + static constexpr float kExpected = 1.0f + 0.5f * kChainLen; + + void SetUp() override { + if (DefaultWebGpuExecutionProvider() == nullptr) { + GTEST_SKIP() << "WebGPU execution provider is not available."; + } + ASSERT_NO_FATAL_FAILURE(BuildAddChainModel(kChainLen, kNumElements, model_bytes_)); + keepalive_ = MakeSession(); + } + + std::unique_ptr MakeSession() { + SessionOptions so; + so.session_logid = "webgpu_concurrent_ctx"; + auto session = std::make_unique(so, GetEnvironment()); + ORT_THROW_IF_ERROR(session->RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + ORT_THROW_IF_ERROR(session->Load(model_bytes_.data(), static_cast(model_bytes_.size()))); + ORT_THROW_IF_ERROR(session->Initialize()); + return session; + } + + NameMLValMap MakeFeeds() const { + std::vector x_values(static_cast(kNumElements), 1.0f); + OrtValue x_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], + std::vector{kNumElements}, x_values, &x_value); + return NameMLValMap{{"X", x_value}}; + } + + // Runs one inference and validates the numerical result. Records into `sink` on failure. + void RunOnce(InferenceSession& session, ErrorSink& sink, const std::string& tag) { + std::vector output_names{"Y"}; + std::vector fetches; + Status s = session.Run(RunOptions{}, MakeFeeds(), output_names, &fetches); + if (!s.IsOK()) { + sink.Record(tag + " Run failed: " + s.ErrorMessage()); + return; + } + const Tensor& out = fetches[0].Get(); + const float* data = out.Data(); + const int64_t n = out.Shape().Size(); + for (int64_t i = 0; i < n; i += (n / 8) + 1) { + if (std::abs(data[i] - kExpected) > 1e-3f) { + sink.Record(tag + " wrong output: " + std::to_string(data[i])); + return; + } + } + } + + // Repeatedly runs an existing session until `iters` reached or a failure is recorded. + void RunLoop(InferenceSession& session, int iters, ErrorSink& sink, const std::string& tag) { + for (int i = 0; i < iters && !sink.Failed(); ++i) { + try { + RunOnce(session, sink, tag); + } catch (const std::exception& e) { + sink.Record(tag + " threw: " + e.what()); + return; + } + } + } + + static void JoinAll(std::vector& threads) { + for (auto& t : threads) { + t.join(); + } + } + + std::string model_bytes_; + std::unique_ptr keepalive_; +}; + +// Case A: one session, run() concurrently from many threads. InferenceSession serializes these +// via session_mutex_, but this still must never crash or deadlock on the shared context. +TEST_F(WebGpuConcurrentContextTest, SingleSessionMultiThreadRun) { + constexpr int kThreads = 4; + constexpr int kIters = 30; + + auto session = MakeSession(); + ErrorSink sink; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { RunLoop(*session, kIters, sink, "A.run" + std::to_string(t)); }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + +// Case B: many threads, each with its own pre-created session. All sessions share the default +// context, so their Run paths execute concurrently against it. Before the fix they also shared a +// single command encoder, which is what this case exposed. +TEST_F(WebGpuConcurrentContextTest, PerThreadSessionRun) { + constexpr int kThreads = 4; + constexpr int kIters = 30; + + std::vector> sessions(kThreads); + for (int t = 0; t < kThreads; ++t) { + sessions[t] = MakeSession(); + } + + ErrorSink sink; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { RunLoop(*sessions[t], kIters, sink, "B.sess" + std::to_string(t)); }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + +// Case C: mixed. Runner threads dispatch on pre-initialized sessions while builder threads keep +// creating + Initializing (initializer-upload path) + running fresh sessions. This is the +// closest shape to the original WebNN crash (dispatch racing initializer upload). +TEST_F(WebGpuConcurrentContextTest, MixedCreateAndRun) { + constexpr int kRunners = 3; + constexpr int kBuilders = 3; + constexpr int kIters = 30; + + std::vector> runner_sessions(kRunners); + for (int t = 0; t < kRunners; ++t) { + runner_sessions[t] = MakeSession(); + } + + ErrorSink sink; + std::vector threads; + for (int t = 0; t < kRunners; ++t) { + threads.emplace_back([&, t]() { RunLoop(*runner_sessions[t], kIters, sink, "C.runner" + std::to_string(t)); }); + } + for (int t = 0; t < kBuilders; ++t) { + threads.emplace_back([&, t]() { + const std::string tag = "C.builder" + std::to_string(t); + for (int i = 0; i < kIters && !sink.Failed(); ++i) { + try { + auto session = MakeSession(); // Initialize -> initializer upload + RunOnce(*session, sink, tag); + } catch (const std::exception& e) { + sink.Record(tag + " threw: " + e.what()); + return; + } + } + }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + +// Case D: churn. Many threads each repeatedly create + run + destroy their own session, +// exercising concurrent allocation and release against the shared context-level buffer manager. +TEST_F(WebGpuConcurrentContextTest, ChurnCreateRunDestroy) { + constexpr int kThreads = 4; + constexpr int kIters = 15; + + ErrorSink sink; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { + const std::string tag = "D.churn" + std::to_string(t); + for (int i = 0; i < kIters && !sink.Failed(); ++i) { + try { + auto session = MakeSession(); + RunOnce(*session, sink, tag); + session.reset(); // drop session -> release its resources while others run + } catch (const std::exception& e) { + sink.Record(tag + " threw: " + e.what()); + return; + } + } + }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + +// Case E: a session warming up pipelines concurrently with a warm session must remain correct. +// +// A builder thread creates a session over a model with many distinct op types and runs it once, +// which compiles one pipeline per op. Meanwhile a runner thread keeps dispatching on an already +// warm session. +// +// Timing is reported to make lock contention visible without imposing a machine-dependent +// performance threshold on this correctness test. +TEST_F(WebGpuConcurrentContextTest, ColdAndWarmSessionsRunConcurrently) { + using Clock = std::chrono::steady_clock; + using Ms = std::chrono::duration; + + std::string cold_model_bytes; + ASSERT_NO_FATAL_FAILURE(BuildUnaryFanOutModel(kNumElements, cold_model_bytes)); + + // Warm session: every pipeline it needs is already compiled and cached on the context. + auto warm_session = MakeSession(); + ErrorSink sink; + RunOnce(*warm_session, sink, "E.warmup"); + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); + + // Baseline latency with no other session active. + constexpr int kBaselineIters = 20; + std::vector baseline_ms; + baseline_ms.reserve(kBaselineIters); + for (int i = 0; i < kBaselineIters; ++i) { + const auto start = Clock::now(); + RunOnce(*warm_session, sink, "E.baseline"); + baseline_ms.push_back(Ms(Clock::now() - start).count()); + } + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); + std::sort(baseline_ms.begin(), baseline_ms.end()); + const double baseline_median = baseline_ms[baseline_ms.size() / 2]; + + std::atomic builder_done{false}; + std::vector contended_ms; + double builder_ms = 0.0; + + std::thread builder([&]() { + const auto start = Clock::now(); + try { + SessionOptions so; + so.session_logid = "webgpu_concurrent_ctx_cold"; + InferenceSession cold_session(so, GetEnvironment()); + ORT_THROW_IF_ERROR(cold_session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + ORT_THROW_IF_ERROR(cold_session.Load(cold_model_bytes.data(), static_cast(cold_model_bytes.size()))); + ORT_THROW_IF_ERROR(cold_session.Initialize()); + + std::vector output_names; + for (size_t i = 0; i < std::size(kUnaryOps); ++i) { + output_names.push_back("Y" + std::to_string(i)); + } + std::vector fetches; + // This Run is what triggers compilation of one pipeline per op type. + ORT_THROW_IF_ERROR(cold_session.Run(RunOptions{}, MakeFeeds(), output_names, &fetches)); + } catch (const std::exception& e) { + sink.Record(std::string("E.builder threw: ") + e.what()); + } + builder_ms = Ms(Clock::now() - start).count(); + builder_done.store(true); + }); + + // Always take at least one measurement so the loop cannot produce an empty sample set if the + // builder happens to finish first. + const auto contended_start = Clock::now(); + do { + const auto start = Clock::now(); + RunOnce(*warm_session, sink, "E.contended"); + contended_ms.push_back(Ms(Clock::now() - start).count()); + } while (!builder_done.load() && !sink.Failed()); + const double contended_window_ms = Ms(Clock::now() - contended_start).count(); + builder.join(); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); + + const double contended_max = *std::max_element(contended_ms.begin(), contended_ms.end()); + std::sort(contended_ms.begin(), contended_ms.end()); + const double contended_median = contended_ms[contended_ms.size() / 2]; + + // Fraction of the contended window during which the warm session kept doing useful work, using + // its uncontended latency as the reference. 1.0 means the cold session cost it nothing. + const double throughput_efficiency = + (static_cast(contended_ms.size()) * baseline_median) / contended_window_ms; + + std::cout << "[ WebGPU ] cold-session warmup: " << builder_ms << " ms for " + << std::size(kUnaryOps) << " pipelines\n" + << "[ WebGPU ] warm-session latency: baseline median " << baseline_median + << " ms, contended median " << contended_median + << " ms, contended max " << contended_max + << " ms over " << contended_ms.size() << " runs\n" + << "[ WebGPU ] warm-session throughput efficiency while cold session warmed up: " + << throughput_efficiency << std::endl; +} + +// Case F: the public session allocator wraps this same internal allocator. Allocations made +// concurrently with Run must not race the session's command recording state. +TEST_F(WebGpuConcurrentContextTest, SessionAllocatorAndRunConcurrently) { + constexpr int kIters = 40; + auto session = MakeSession(); + auto allocator = session->GetAllocator(OrtMemoryInfo(WEBGPU_BUFFER, + OrtAllocatorType::OrtDeviceAllocator, + webgpu::WebGpuDevice, + OrtMemTypeDefault)); + ASSERT_NE(allocator, nullptr); + + ErrorSink sink; + std::barrier start{2}; + std::thread runner([&]() { + start.arrive_and_wait(); + RunLoop(*session, kIters, sink, "F.run"); + }); + std::thread allocator_user([&]() { + start.arrive_and_wait(); + try { + for (int i = 0; i < kIters && !sink.Failed(); ++i) { + Tensor tensor(DataTypeImpl::GetType(), TensorShape{kNumElements}, allocator); + ASSERT_NE(tensor.MutableDataRaw(), nullptr); + } + } catch (const std::exception& e) { + sink.Record(std::string("F.allocator threw: ") + e.what()); + } + }); + runner.join(); + allocator_user.join(); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + +// Case G: exercise the plugin Env allocator's shared implementation and getter ownership in a +// native build. Its private command state must remain thread-safe across callers. +TEST_F(WebGpuConcurrentContextTest, SharedAllocatorMultiThreadCreateTensor) { + constexpr int kThreads = 4; + constexpr int kIters = 60; + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + auto context_ref = std::shared_ptr(&context, [](webgpu::WebGpuContext*) {}); + auto recording = std::make_shared(); + std::weak_ptr recording_lifetime = recording; + auto allocator = std::make_shared( + [context_ref = std::move(context_ref)]() -> const webgpu::BufferManager& { + return context_ref->BufferManager(); + }, + [recording = std::move(recording)]() -> webgpu::CommandRecordingState& { return *recording; }, + false, + []() { return true; }); + + ErrorSink sink; + std::barrier start{kThreads}; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { + start.arrive_and_wait(); + try { + for (int i = 0; i < kIters && !sink.Failed(); ++i) { + Tensor tensor(DataTypeImpl::GetType(), TensorShape{1024 + t * 16}, allocator); + ASSERT_NE(tensor.MutableDataRaw(), nullptr); + } + } catch (const std::exception& e) { + sink.Record("G.thread" + std::to_string(t) + " threw: " + e.what()); + } + }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); + EXPECT_FALSE(recording_lifetime.expired()); + AllocatorStats stats{}; + allocator->GetStats(&stats); + EXPECT_EQ(stats.num_allocs, 0); + allocator.reset(); + EXPECT_TRUE(recording_lifetime.expired()); +} + +// Case H: OrtEnv owns one data-transfer implementation per EP factory. Concurrent CopyTensors +// calls must not encode and flush through the same recording state simultaneously. +TEST_F(WebGpuConcurrentContextTest, SharedDataTransferMultiThreadCopy) { + constexpr int kThreads = 4; + constexpr int kIters = 30; + constexpr size_t kElements = 4096; + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + webgpu::CommandRecordingState recording; + webgpu::DataTransferImpl data_transfer( + [&context]() -> const webgpu::BufferManager& { return context.BufferManager(); }, recording); + + std::array gpu_buffers; + for (auto& buffer : gpu_buffers) { + wgpu::BufferDescriptor desc{}; + desc.size = kElements * sizeof(float); + desc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + buffer = context.Device().CreateBuffer(&desc); + } + + ErrorSink sink; + std::barrier start{kThreads}; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { + std::vector input(kElements, static_cast(t + 1)); + std::vector output(kElements); + start.arrive_and_wait(); + try { + for (int i = 0; i < kIters && !sink.Failed(); ++i) { + ORT_THROW_IF_ERROR(data_transfer.CopyTensor(input.data(), false, gpu_buffers[t].Get(), true, + input.size() * sizeof(float))); + ORT_THROW_IF_ERROR(data_transfer.CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, + output.size() * sizeof(float))); + if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { + sink.Record("H.thread" + std::to_string(t) + " copied incorrect data"); + } + } + } catch (const std::exception& e) { + sink.Record("H.thread" + std::to_string(t) + " threw: " + e.what()); + } + }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + +// Case I: separate data-transfer objects share the context-level BufferManager but own distinct +// command recording timelines. Their command encoders and pending buffers must remain isolated. +TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { + constexpr int kThreads = 4; + constexpr int kIters = 30; + constexpr size_t kElements = 4096; + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + + std::array recordings; + std::array, kThreads> data_transfers; + std::array gpu_buffers; + for (int t = 0; t < kThreads; ++t) { + data_transfers[t] = std::make_unique( + [&context]() -> const webgpu::BufferManager& { return context.BufferManager(); }, recordings[t]); + wgpu::BufferDescriptor desc{}; + desc.size = kElements * sizeof(float); + desc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + gpu_buffers[t] = context.Device().CreateBuffer(&desc); + } + + ErrorSink sink; + std::barrier start{kThreads}; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { + std::vector input(kElements, static_cast(t + 1)); + std::vector output(kElements); + start.arrive_and_wait(); + try { + for (int i = 0; i < kIters && !sink.Failed(); ++i) { + ORT_THROW_IF_ERROR(data_transfers[t]->CopyTensor(input.data(), false, gpu_buffers[t].Get(), true, + input.size() * sizeof(float))); + ORT_THROW_IF_ERROR(data_transfers[t]->CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, + output.size() * sizeof(float))); + if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { + sink.Record("I.thread" + std::to_string(t) + " copied incorrect data"); + } + } + } catch (const std::exception& e) { + sink.Record("I.thread" + std::to_string(t) + " threw: " + e.what()); + } + }); + } + JoinAll(threads); + + ASSERT_FALSE(sink.Failed()) << sink.FirstError(); +} + // DIAGNOSTIC (disabled by default): keeps N sessions over an identical model alive and runs them // round-robin on one thread, then holds so an external sampler can read steady-state GPU memory. // This is the shape that distinguishes a context-wide buffer pool (steady state ~= one session's diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 0b9473ae00257..d7c1610cf6ab5 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -5,6 +5,8 @@ #include #include #include +#include +#include #include #include #include @@ -20,14 +22,26 @@ #include "core/common/common.h" #include "core/framework/config_options.h" #include "core/framework/run_options.h" +#include "core/framework/tensor.h" +#include "core/graph/onnx_protobuf.h" #include "core/providers/webgpu/allocator.h" #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/webgpu_context.h" #include "core/providers/webgpu/webgpu_execution_provider.h" +#include "core/providers/webgpu/webgpu_execution_provider.h" #include "core/providers/webgpu/webgpu_provider_factory_creator.h" #include "core/providers/webgpu/webgpu_provider_options.h" #include "core/session/onnxruntime_session_options_config_keys.h" +#include "core/session/inference_session.h" +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" #include "test/util/include/asserts.h" +#include "test/util/include/default_providers.h" +#include "test/util/include/temp_dir.h" + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +#include "core/providers/webgpu/d3d12_external_data_loader.h" +#endif #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) #include "dawn/native/DawnNative.h" @@ -80,7 +94,7 @@ void RunWithFreshDefaultContext(TestBody test_body, bool compile_only_parent = f } existing_provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); ASSERT_NE(existing_provider, nullptr); - existing_device = webgpu::WebGpuContextFactory::GetContext(0).Device().Get(); + existing_device = webgpu::WebGpuContextFactory::GetContext(existing_provider->GetDeviceId()).Device().Get(); } EXPECT_EXIT( @@ -90,7 +104,8 @@ void RunWithFreshDefaultContext(TestBody test_body, bool compile_only_parent = f }, testing::ExitedWithCode(EXIT_SUCCESS), ""); - EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(0).Device().Get(), existing_device); + EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(existing_provider->GetDeviceId()).Device().Get(), + existing_device); #if defined(GTEST_HAS_ABSL) && !defined(GTEST_NO_ABSL_FLAGS) flag_error.clear(); EXPECT_TRUE(death_test_style_flag->ParseFrom(previous_style, &flag_error)) << flag_error; @@ -103,6 +118,18 @@ void RunWithFreshDefaultContext(TestBody test_body, bool compile_only_parent = f #endif } +ConfigOptions WeightLoadAccelerationOptions(const char* value) { + ConfigOptions options; + ORT_THROW_IF_ERROR(options.AddConfigEntry(kWeightLoadAcceleration, value)); + return options; +} + +ConfigOptions CompileOnlyWeightLoadAccelerationOptions(const char* value) { + auto options = WeightLoadAccelerationOptions(value); + ORT_THROW_IF_ERROR(options.AddConfigEntry(kOrtSessionOptionCompileOnly, "1")); + return options; +} + bool DeviceToggleIsEnabled(const webgpu::WebGpuContext& context, std::string_view toggle_name) { #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) const auto toggles = dawn::native::GetTogglesUsed(context.Device().Get()); @@ -120,9 +147,11 @@ bool DisableRobustnessToggleIsEnabled(const webgpu::WebGpuContext& context) { return DeviceToggleIsEnabled(context, "disable_robustness"); } -std::array ReadBufferWithExternalCommandEncoder(webgpu::WebGpuContext& context, - WGPUBuffer buffer) { - constexpr size_t kBufferSize = sizeof(std::array); +template +std::array ReadBufferWithExternalCommandEncoder( + webgpu::WebGpuContext& context, WGPUBuffer buffer) { + constexpr size_t kBufferSize = + sizeof(std::array); wgpu::BufferDescriptor readback_desc{}; readback_desc.size = kBufferSize; readback_desc.usage = wgpu::BufferUsage::MapRead | wgpu::BufferUsage::CopyDst; @@ -142,7 +171,7 @@ std::array ReadBufferWithExternalCommandEncoder(webgpu::WebGpuCont &map_status))); ORT_ENFORCE(map_status == wgpu::MapAsyncStatus::Success); - std::array result; + std::array result; const auto* mapped_data = static_cast(readback_buffer.GetConstMappedRange()); ORT_ENFORCE(mapped_data != nullptr); std::copy_n(mapped_data, result.size(), result.begin()); @@ -1068,7 +1097,7 @@ TEST(WebGpuContextTest, AdapterIndexAcceptsNonNegativeInteger) { auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); ASSERT_NE(ep, nullptr); - EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(0).Device().Get(), nullptr); + EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()).Device().Get(), nullptr); }, /*compile_only_parent=*/true); #endif @@ -1108,7 +1137,6 @@ TEST(WebGpuContextTest, AdapterIndexRejectsUnsupportedBuild) { #else ConfigOptions options; ORT_THROW_IF_ERROR(options.AddConfigEntry(kAdapterIndex, "0")); - ORT_THROW_IF_ERROR(options.AddConfigEntry(kOrtSessionOptionCompileOnly, "1")); try { WebGpuProviderFactoryCreator::Create(options); @@ -1119,6 +1147,574 @@ TEST(WebGpuContextTest, AdapterIndexRejectsUnsupportedBuild) { #endif } +TEST(WebGpuContextTest, WeightLoadAccelerationRejectsBooleanAndUnknownModes) { + EXPECT_THROW(WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions("1")), + OnnxRuntimeException); + EXPECT_THROW(WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions("automatic")), + OnnxRuntimeException); +} + +TEST(WebGpuContextTest, WeightLoadAccelerationOffDoesNotRequireDeviceSupport) { + auto ep = WebGpuProviderFactoryCreator::Create( + CompileOnlyWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Off)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); +} + +TEST(WebGpuContextTest, PreferredWeightLoadAccelerationFallsBackWithoutDeviceSupport) { + auto ep = WebGpuProviderFactoryCreator::Create( + CompileOnlyWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Preferred)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); +} + +TEST(WebGpuContextTest, RequiredWeightLoadAccelerationFailsWithoutDeviceSupport) { + auto factory = WebGpuProviderFactoryCreator::Create( + CompileOnlyWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Required)); +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + auto ep = factory->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + EXPECT_FALSE(loader->BeginLoad().IsOK()); +#else + EXPECT_THROW(factory->CreateProvider(), OnnxRuntimeException); +#endif +} + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +TEST(WebGpuContextTest, D3D12AcceleratedCanBeEnabledAfterInitialOffSession) { + auto off_factory = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Off)); + auto off_ep = off_factory->CreateProvider(); + ASSERT_NE(off_ep, nullptr); + + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + const auto support_status = + webgpu::CheckD3D12AcceleratedExternalWeightsSupport(context); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + auto required_factory = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Required)); + auto required_ep = required_factory->CreateProvider(); + ASSERT_NE(required_ep, nullptr); + auto loader = required_ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + EXPECT_STATUS_OK(loader->BeginLoad()); + loader->AbortLoad(); +} + +TEST(WebGpuContextTest, D3D12AcceleratedAllocatorUsesProviderRecording) { + auto ep = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions( + kWeightLoadAcceleration_Preferred)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto allocators = ep->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + auto& allocator = allocators.front(); + void* buffer = allocator->Alloc(16); + ASSERT_NE(buffer, nullptr); + + auto& webgpu_ep = *static_cast(ep.get()); + auto& recording = webgpu_ep.Recording(); + recording.has_unsubmitted_work = true; + allocator->Free(buffer); + EXPECT_EQ(recording.pending_buffers.size(), 1u); + + recording.has_unsubmitted_work = false; + webgpu_ep.InitializerBufferManager().RefreshPendingBuffers(recording); + EXPECT_TRUE(recording.pending_buffers.empty()); +} + +TEST(WebGpuContextTest, DuplicateProviderDoesNotRegisterD3D12AcceleratedLoader) { + InferenceSession session{SessionOptions{}, GetEnvironment()}; + auto off_provider = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions( + kWeightLoadAcceleration_Off)) + ->CreateProvider(); + ASSERT_NE(off_provider, nullptr); + auto allocators = off_provider->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto memory_info = allocators.front()->Info(); + ASSERT_STATUS_OK( + session.RegisterExecutionProvider(std::move(off_provider))); + + auto accelerated_provider = + WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Required)) + ->CreateProvider(); + ASSERT_NE(accelerated_provider, nullptr); + const auto status = + session.RegisterExecutionProvider(std::move(accelerated_provider)); + EXPECT_FALSE(status.IsOK()); + EXPECT_NE(status.ErrorMessage().find("already been registered"), + std::string::npos); + EXPECT_EQ(session.GetExternalDataLoaderManager().GetExternalDataLoader( + memory_info), + nullptr); +} + +TEST(WebGpuContextTest, D3D12AcceleratedLoadsExternalTensorsAcrossFilesAndRanges) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_accelerated_external_tensor_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + const auto other_data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("other_weights.bin"); + constexpr size_t kDataOffset = 32; + constexpr size_t kSeparatedDataOffset = 128; + const std::array expected{ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14}; + const std::array separated_expected{101, 102, 103, 104}; + const std::array other_expected{201, 202, 203}; + { + std::ofstream stream{data_path, + std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + const std::array prefix{}; + stream.write(prefix.data(), static_cast(prefix.size())); + stream.write(reinterpret_cast(expected.data()), + static_cast(sizeof(expected))); + stream.seekp(kSeparatedDataOffset); + stream.write(reinterpret_cast(separated_expected.data()), + static_cast(sizeof(separated_expected))); + ASSERT_TRUE(stream.good()); + } + { + std::ofstream stream{other_data_path, + std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + const std::array prefix{}; + stream.write(prefix.data(), static_cast(prefix.size())); + stream.write(reinterpret_cast(other_expected.data()), + static_cast(sizeof(other_expected))); + ASSERT_TRUE(stream.good()); + } + + auto options = + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Required); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + EXPECT_FALSE(loader->SupportsDataType( + ONNX_NAMESPACE::TensorProto_DataType_BOOL)); + EXPECT_TRUE(loader->SupportsDataType( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + auto allocators = ep->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto& allocator = allocators.front(); + + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "weights", kDataOffset, sizeof(expected))); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "separated", kSeparatedDataOffset, + sizeof(separated_expected))); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), other_data_path, "other_weights", kDataOffset, + sizeof(other_expected))); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "empty", + kSeparatedDataOffset + sizeof(separated_expected), 0)); + ASSERT_STATUS_OK(loader->FinalizeLoad([]() { return false; })); + + Tensor weights{DataTypeImpl::GetType(), TensorShape({15}), + nullptr, allocator}; + ASSERT_STATUS_OK(loader->LoadTensor( + Env::Default(), data_path, "weights", kDataOffset, sizeof(expected), + allocator, weights)); + std::array aligned_expected{}; + std::copy(expected.begin(), expected.end(), aligned_expected.begin()); + EXPECT_EQ(ReadBufferWithExternalCommandEncoder<16>( + webgpu::WebGpuContextFactory::GetContext(0), + reinterpret_cast(weights.MutableDataRaw())), + aligned_expected); + + Tensor separated_weights{DataTypeImpl::GetType(), TensorShape({4}), + nullptr, allocator}; + ASSERT_STATUS_OK(loader->LoadTensor( + Env::Default(), data_path, "separated", kSeparatedDataOffset, + sizeof(separated_expected), allocator, separated_weights)); + EXPECT_EQ(ReadBufferWithExternalCommandEncoder<4>( + webgpu::WebGpuContextFactory::GetContext(0), + reinterpret_cast( + separated_weights.MutableDataRaw())), + separated_expected); + + Tensor other_weights{DataTypeImpl::GetType(), TensorShape({3}), + nullptr, allocator}; + ASSERT_STATUS_OK(loader->LoadTensor( + Env::Default(), other_data_path, "other_weights", kDataOffset, + sizeof(other_expected), allocator, other_weights)); + std::array aligned_other_expected{}; + std::copy(other_expected.begin(), other_expected.end(), + aligned_other_expected.begin()); + aligned_other_expected.back() = other_expected.front(); + EXPECT_EQ(ReadBufferWithExternalCommandEncoder<4>( + webgpu::WebGpuContextFactory::GetContext(0), + reinterpret_cast(other_weights.MutableDataRaw())), + aligned_other_expected); + + Tensor empty{DataTypeImpl::GetType(), TensorShape({0}), nullptr, + allocator}; + ASSERT_STATUS_OK(loader->LoadTensor( + Env::Default(), data_path, "empty", + kSeparatedDataOffset + sizeof(separated_expected), 0, allocator, empty)); + EXPECT_EQ(empty.SizeInBytes(), 0u); + EXPECT_EQ(empty.MutableDataRaw(), nullptr); +} + +void RunD3D12AcceleratedExternalInitializerSessionTest( + const char* mode, bool fail_late_registration = false) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_accelerated_session_test")}; + const auto model_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("model.onnx"); + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + constexpr size_t kDataOffset = 32; + const std::array weights{ + 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, + 10.0f, 11.0f, 12.0f, 13.0f, 14.0f}; + { + std::ofstream stream{data_path, std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + const std::array prefix{}; + stream.write(prefix.data(), static_cast(prefix.size())); + stream.write(reinterpret_cast(weights.data()), + static_cast(sizeof(weights))); + ASSERT_TRUE(stream.good()); + } + + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(8); + model.set_producer_name("onnxruntime-test"); + auto* opset = model.add_opset_import(); + opset->set_domain(""); + opset->set_version(13); + auto* graph = model.mutable_graph(); + graph->set_name("webgpu_d3d12_accelerated_session"); + + const auto set_tensor_type = [&weights]( + ONNX_NAMESPACE::ValueInfoProto* value_info, + const char* name) { + value_info->set_name(name); + auto* tensor_type = value_info->mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor_type->mutable_shape()->add_dim()->set_dim_value(weights.size()); + }; + set_tensor_type(graph->add_input(), "X"); + set_tensor_type(graph->add_output(), "Y"); + + auto* initializer = graph->add_initializer(); + initializer->set_name("W"); + initializer->set_data_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + initializer->add_dims(weights.size()); + initializer->set_data_location( + ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + auto* location = initializer->add_external_data(); + location->set_key("location"); + location->set_value("weights.bin"); + auto* offset = initializer->add_external_data(); + offset->set_key("offset"); + offset->set_value(std::to_string(kDataOffset)); + auto* length = initializer->add_external_data(); + length->set_key("length"); + length->set_value(std::to_string(sizeof(weights))); + + auto* add = graph->add_node(); + add->set_op_type("Add"); + add->add_input("X"); + add->add_input("W"); + add->add_output("Y"); + { + std::ofstream stream{model_path, std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + ASSERT_TRUE(model.SerializeToOstream(&stream)); + } + + SessionOptions session_options; + InferenceSession session{session_options, GetEnvironment()}; + auto provider = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions(mode)) + ->CreateProvider(); + ASSERT_NE(provider, nullptr); + if (fail_late_registration) { + auto allocators = provider->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto memory_info = allocators.front()->Info(); + ASSERT_STATUS_OK(session.Load(model_path.native())); + ASSERT_TRUE(std::filesystem::remove(data_path)); + const auto status = + session.RegisterExecutionProvider(std::move(provider)); + EXPECT_FALSE(status.IsOK()); + EXPECT_EQ(session.GetExternalDataLoaderManager().GetExternalDataLoader( + memory_info), + nullptr); + return; + } + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(provider))); + ASSERT_STATUS_OK(session.Load(model_path.native())); + ASSERT_STATUS_OK(session.Initialize()); + + std::vector input(weights.size(), 1.0f); + OrtValue input_value; + CreateMLValue( + TestCPUExecutionProvider()->CreatePreferredAllocators()[0], + {static_cast(weights.size())}, input, &input_value); + NameMLValMap feeds{{"X", input_value}}; + std::vector output_names{"Y"}; + std::vector fetches; + ASSERT_STATUS_OK( + session.Run(RunOptions{}, feeds, output_names, &fetches)); + ASSERT_EQ(fetches.size(), 1u); + const auto& output = fetches[0].Get(); + ASSERT_EQ(output.Shape(), TensorShape({static_cast(weights.size())})); + const auto* output_data = output.Data(); + for (size_t index = 0; index < weights.size(); ++index) { + EXPECT_FLOAT_EQ(output_data[index], input[index] + weights[index]); + } +} + +TEST(WebGpuContextTest, D3D12AcceleratedLoadsExternalInitializerThroughSession) { + auto probe_provider = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions( + kWeightLoadAcceleration_Off)) + ->CreateProvider(); + ASSERT_NE(probe_provider, nullptr); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + RunD3D12AcceleratedExternalInitializerSessionTest( + kWeightLoadAcceleration_Required); +} + +TEST(WebGpuContextTest, LateRegistrationFailureRollsBackD3D12AcceleratedLoader) { + auto probe_provider = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions( + kWeightLoadAcceleration_Off)) + ->CreateProvider(); + ASSERT_NE(probe_provider, nullptr); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + RunD3D12AcceleratedExternalInitializerSessionTest( + kWeightLoadAcceleration_Required, true); +} + +TEST(WebGpuContextTest, PreferredD3D12AcceleratedPreservesCancellation) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_accelerated_cancellation_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + const std::array data{}; + { + std::ofstream stream{data_path, + std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + stream.write(reinterpret_cast(data.data()), + static_cast(sizeof(data))); + ASSERT_TRUE(stream.good()); + } + + auto options = + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Preferred); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "weights", 0, sizeof(data))); + const auto status = loader->FinalizeLoad([]() { return true; }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->AbortLoad(); +} + +TEST(WebGpuContextTest, D3D12AcceleratedZeroRequestFinalBatchPreservesCancellation) { + auto options = + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Preferred); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + ASSERT_STATUS_OK(loader->BeginLoad()); + const auto status = loader->FinalizeLoad([]() { return true; }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->AbortLoad(); +} + +TEST(WebGpuContextTest, D3D12AcceleratedChecksCancellationDuringImport) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_accelerated_import_cancellation_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + const std::array data{}; + { + std::ofstream stream{data_path, + std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + stream.write(reinterpret_cast(data.data()), + static_cast(sizeof(data))); + ASSERT_TRUE(stream.good()); + } + + auto options = + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Preferred); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "weights", 0, sizeof(data))); + + size_t cancellation_checks = 0; + const auto status = loader->FinalizeLoad([&cancellation_checks]() { + return cancellation_checks++ != 0; + }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->AbortLoad(); +} + +TEST(WebGpuContextTest, PreferredD3D12AcceleratedFallsBackAfterOperationalFailure) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_accelerated_fallback_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + const std::array data{}; + { + std::ofstream stream{data_path, + std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + stream.write(reinterpret_cast(data.data()), + static_cast(sizeof(data))); + ASSERT_TRUE(stream.good()); + } + + auto options = + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Preferred); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + auto allocators = ep->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto& allocator = allocators.front(); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "weights", 0, sizeof(data))); + ASSERT_TRUE(std::filesystem::remove(data_path)); + EXPECT_STATUS_OK(loader->FinalizeLoad([]() { return false; })); + EXPECT_FALSE(loader->CanLoad(allocator->Info())); +} + +TEST(WebGpuContextTest, PreferredD3D12AcceleratedPreservesConcurrentCancellation) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_accelerated_failure_cancellation_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + const std::array data{}; + { + std::ofstream stream{data_path, + std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + stream.write(reinterpret_cast(data.data()), + static_cast(sizeof(data))); + ASSERT_TRUE(stream.good()); + } + + auto options = + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Preferred); + auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( + webgpu::WebGpuContextFactory::GetContext(0)); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "weights", 0, sizeof(data))); + ASSERT_TRUE(std::filesystem::remove(data_path)); + const auto status = loader->FinalizeLoad([]() { return true; }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->AbortLoad(); +} + +TEST(WebGpuContextTest, WeightLoadAccelerationModeResolution) { + const auto unsupported = + ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "D3D12Accelerated unavailable"); + bool enabled = true; + + EXPECT_STATUS_OK(webgpu::ResolveWeightLoadAccelerationMode( + webgpu::WeightLoadAccelerationMode::Off, unsupported, enabled)); + EXPECT_FALSE(enabled); + + enabled = true; + EXPECT_STATUS_OK(webgpu::ResolveWeightLoadAccelerationMode( + webgpu::WeightLoadAccelerationMode::Preferred, unsupported, enabled)); + EXPECT_FALSE(enabled); + + enabled = true; + const auto required_status = webgpu::ResolveWeightLoadAccelerationMode( + webgpu::WeightLoadAccelerationMode::Required, unsupported, + enabled); + EXPECT_FALSE(required_status.IsOK()); + EXPECT_FALSE(enabled); + EXPECT_EQ(required_status.ErrorMessage(), unsupported.ErrorMessage()); +} +#endif + TEST(WebGpuContextTest, AdapterIndexRejectsConflictingSelectorOnReusedContext) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Physical adapter enumeration requires a native Dawn build."; @@ -1150,11 +1746,26 @@ TEST(WebGpuContextTest, CompileOnlyContextDoesNotCreateDevice) { auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); ASSERT_NE(ep, nullptr); - EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(0).Device().Get(), nullptr); + EXPECT_EQ(ep->GetDeviceId(), webgpu::kDeviceFreeDefaultContextId); + EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()).Device().Get(), nullptr); }, /*compile_only_parent=*/true); } +TEST(WebGpuContextTest, CompileOnlyAndRunnableDefaultContextsAreIsolated) { + ConfigOptions compile_only_options; + ORT_THROW_IF_ERROR(compile_only_options.AddConfigEntry(kOrtSessionOptionCompileOnly, "1")); + auto compile_only_ep = WebGpuProviderFactoryCreator::Create(compile_only_options)->CreateProvider(); + ASSERT_NE(compile_only_ep, nullptr); + + ConfigOptions runnable_options; + auto runnable_ep = WebGpuProviderFactoryCreator::Create(runnable_options)->CreateProvider(); + ASSERT_NE(runnable_ep, nullptr); + + EXPECT_TRUE(webgpu::WebGpuContextFactory::GetContext(compile_only_ep->GetDeviceId()).IsDeviceFree()); + EXPECT_FALSE(webgpu::WebGpuContextFactory::GetContext(runnable_ep->GetDeviceId()).IsDeviceFree()); +} + TEST(WebGpuContextTest, EnableRobustnessIsIndependentFromValidationMode) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native toggle inspection is unavailable."; diff --git a/onnxruntime/test/webgpu/delay_load/main.cc b/onnxruntime/test/webgpu/delay_load/main.cc index da223988b2a2e..8340725716573 100644 --- a/onnxruntime/test/webgpu/delay_load/main.cc +++ b/onnxruntime/test/webgpu/delay_load/main.cc @@ -95,7 +95,8 @@ int run() { Ort::SessionOptions session_options; session_options.DisableMemPattern(); - std::unordered_map provider_options; + std::unordered_map provider_options{ + {"weightLoadAcceleration", "preferred"}}; session_options.AppendExecutionProvider("WebGPU", provider_options); Ort::Session session{env, MODEL_DATA, sizeof(MODEL_DATA), session_options}; diff --git a/onnxruntime/test/webgpu/external_dawn/main.cc b/onnxruntime/test/webgpu/external_dawn/main.cc index a3024672b2b71..017e13bec44e0 100644 --- a/onnxruntime/test/webgpu/external_dawn/main.cc +++ b/onnxruntime/test/webgpu/external_dawn/main.cc @@ -169,6 +169,9 @@ int main(int argc, char* argv[]) { std::unordered_map provider_options; if (!no_proc_table) { provider_options["dawnProcTable"] = std::to_string(reinterpret_cast(&dawn::native::GetProcs())); +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + provider_options["weightLoadAcceleration"] = "preferred"; +#endif } const int custom_id = cache_pass == 2 ? std::numeric_limits::max() : static_cast(cache_pass + 1); if (host_device_mode) { diff --git a/setup.py b/setup.py index fcde855cf20b0..58aadc75b5010 100644 --- a/setup.py +++ b/setup.py @@ -468,7 +468,7 @@ def finalize_options(self): # DirectML Libs libs.extend(["DirectML.dll"]) # WebGPU/Dawn Libs - libs.extend(["dxcompiler.dll"]) + libs.extend(["dxcompiler.dll", "dxil.dll"]) # QNN V68/V73/V81 dependencies qnn_deps = [ "QnnCpu.dll", From c10c419da8442aab4d33a8427c9f38c8c5740437 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:40:23 -0700 Subject: [PATCH 02/46] Remove dxil.dll from the list of WebGPU/Dawn libraries in setup.py --- .../d3d12_file_buffer_loader.cc | 1613 +++++++++-------- setup.py | 2 +- 2 files changed, 843 insertions(+), 772 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 62c21f770ae82..1e872f4886bef 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -119,892 +119,963 @@ struct D3D12FileBufferLoader::Impl { Batch batch; }; - Impl(ID3D12Device* device, const Config& config) - : device_(device), config_(config) { - } + Impl(ID3D12Device* device, const Config& config); + ~Impl(); - ~Impl() { - for (auto& slot : slots_) { - if (slot.resource && slot.mapped != nullptr) { - slot.resource->Unmap(0, nullptr); - slot.mapped = nullptr; - } + common::Status Initialize(); + common::Status Load( + const std::vector& ranges, + Batch& result, + const CancellationToken& cancellation); + + private: + common::Status AllocateDestinations( + const std::vector& sizes, + Batch& batch); + common::Status PrepareFiles( + const std::vector& ranges, + const CancellationToken& cancellation, + std::vector& files); + common::Status PrepareSlots( + uint64_t alignment, + const CancellationToken& cancellation); + common::Status IssueRead( + HANDLE file, + UploadSlot& slot, + uint64_t offset, + DWORD size); + common::Status CompleteRead( + HANDLE file, + uint64_t file_size, + UploadSlot& slot, + DWORD& bytes_read); + template + common::Status ReadRegion( + PreparedFile& file, + const FileReadRegion& region, + const std::vector& ranges, + Batch& batch, + bool& allocation_ready, + EnsureAllocationFn& ensure_allocation, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence); + common::Status SubmitCopies( + UploadSlot& slot, + uint64_t chunk_begin, + uint64_t bytes, + const std::vector& range_indices, + const std::vector& ranges, + Batch& batch, + uint64_t& last_submitted_fence); + common::Status TransitionToCommon( + Batch& batch, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence); + common::Status WaitForFence( + uint64_t value, + const CancellationToken& cancellation); + void WaitForFenceUncancelled(uint64_t value) noexcept; + void DrainActiveReads(HANDLE file) noexcept; + + ComPtr device_; + Config config_; + ComPtr copy_queue_; + ComPtr copy_fence_; + wil::unique_handle fence_event_; + std::vector slots_; + uint64_t next_fence_value_ = 0; + std::mutex load_mutex_; +}; + +D3D12FileBufferLoader::Impl::Impl( + ID3D12Device* device, const Config& config) + : device_(device), config_(config) { +} + +D3D12FileBufferLoader::Impl::~Impl() { + for (auto& slot : slots_) { + if (slot.resource && slot.mapped != nullptr) { + slot.resource->Unmap(0, nullptr); + slot.mapped = nullptr; } } +} - common::Status Initialize() { - if (device_ == nullptr) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "D3D12FileBufferLoader requires a non-null D3D12 device."); +common::Status D3D12FileBufferLoader::Impl::Initialize() { + if (device_ == nullptr) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "D3D12FileBufferLoader requires a non-null D3D12 device."); + } + if (config_.upload_slot_count < 2 || + config_.upload_slot_count > MAXIMUM_WAIT_OBJECTS) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "upload_slot_count must be between 2 and ", MAXIMUM_WAIT_OBJECTS, "."); + } + if (config_.upload_slot_size == 0 || + config_.upload_slot_size > std::numeric_limits::max() || + config_.upload_slot_size % kBufferAlignment != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "upload_slot_size must be a non-zero 64 KiB multiple no larger than DWORD_MAX."); + } + if (config_.max_destination_heap_size == 0 || + config_.max_destination_heap_size > 512ull * 1024ull * 1024ull || + config_.max_destination_heap_size % kBufferAlignment != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "max_destination_heap_size must be a non-zero 64 KiB multiple no larger than 512 MiB."); + } + + D3D12_COMMAND_QUEUE_DESC queue_description{}; + queue_description.Type = D3D12_COMMAND_LIST_TYPE_COPY; + HRESULT hr = device_->CreateCommandQueue( + &queue_description, IID_PPV_ARGS(©_queue_)); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateCommandQueue", hr); + } + + hr = device_->CreateFence( + 0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(©_fence_)); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateFence", hr); + } + + fence_event_.reset(CreateEventExW( + nullptr, nullptr, 0, EVENT_MODIFY_STATE | SYNCHRONIZE)); + if (!fence_event_) { + return Win32Error("CreateEventExW(copy fence)", GetLastError()); + } + + const auto upload_heap = HeapProperties(D3D12_HEAP_TYPE_UPLOAD); + const auto upload_description = + BufferDescription(config_.upload_slot_size, D3D12_RESOURCE_FLAG_NONE); + slots_.resize(config_.upload_slot_count); + for (auto& slot : slots_) { + hr = device_->CreateCommittedResource( + &upload_heap, D3D12_HEAP_FLAG_NONE, &upload_description, + D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, + IID_PPV_ARGS(&slot.resource)); + if (FAILED(hr)) { + return HResultError( + "ID3D12Device::CreateCommittedResource(upload)", hr); + } + + hr = slot.resource->Map(0, nullptr, &slot.mapped); + if (FAILED(hr)) { + return HResultError("ID3D12Resource::Map(upload)", hr); + } + + hr = device_->CreateCommandAllocator( + D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&slot.allocator)); + if (FAILED(hr)) { + return HResultError( + "ID3D12Device::CreateCommandAllocator", hr); } - if (config_.upload_slot_count < 2 || - config_.upload_slot_count > MAXIMUM_WAIT_OBJECTS) { + + hr = device_->CreateCommandList( + 0, D3D12_COMMAND_LIST_TYPE_COPY, slot.allocator.Get(), nullptr, + IID_PPV_ARGS(&slot.command_list)); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateCommandList", hr); + } + hr = slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close(initial)", hr); + } + + slot.read_event.reset(CreateEventExW( + nullptr, nullptr, CREATE_EVENT_MANUAL_RESET, + EVENT_MODIFY_STATE | SYNCHRONIZE)); + if (!slot.read_event) { + return Win32Error("CreateEventExW(read)", GetLastError()); + } + } + + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::Impl::Load( + const std::vector& ranges, + Batch& result, + const CancellationToken& cancellation) { + std::lock_guard lock(load_mutex_); + result.Clear(); + + if (ranges.empty()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "At least one file range is required."); + } + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + + std::vector sizes; + sizes.reserve(ranges.size()); + for (size_t index = 0; index < ranges.size(); ++index) { + const auto& range = ranges[index]; + if (range.path.empty()) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, - "upload_slot_count must be between 2 and ", MAXIMUM_WAIT_OBJECTS, "."); + "File range ", index, " has an empty path."); } - if (config_.upload_slot_size == 0 || - config_.upload_slot_size > std::numeric_limits::max() || - config_.upload_slot_size % kBufferAlignment != 0) { + if (range.length == 0) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, - "upload_slot_size must be a non-zero 64 KiB multiple no larger than DWORD_MAX."); + "File range ", index, " has zero length."); } - if (config_.max_destination_heap_size == 0 || - config_.max_destination_heap_size > 512ull * 1024ull * 1024ull || - config_.max_destination_heap_size % kBufferAlignment != 0) { + if (range.offset > + std::numeric_limits::max() - range.length) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, - "max_destination_heap_size must be a non-zero 64 KiB multiple no larger than 512 MiB."); + "File range ", index, " overflows uint64_t."); } + sizes.push_back(range.length); + } - D3D12_COMMAND_QUEUE_DESC queue_description{}; - queue_description.Type = D3D12_COMMAND_LIST_TYPE_COPY; - HRESULT hr = device_->CreateCommandQueue( - &queue_description, IID_PPV_ARGS(©_queue_)); - if (FAILED(hr)) { - return HResultError("ID3D12Device::CreateCommandQueue", hr); + // Destination allocation can run while the CPU opens files and discovers + // the sector alignment required for unbuffered reads. + std::future allocation_future = std::async( + std::launch::async, [this, sizes = std::move(sizes)]() noexcept { + AllocationResult allocation; + try { + allocation.status = + AllocateDestinations(sizes, allocation.batch); + } catch (const std::exception& ex) { + allocation.status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "Destination allocation failed: ", ex.what()); + } catch (...) { + allocation.status = ORT_MAKE_STATUS( + ONNXRUNTIME, RUNTIME_EXCEPTION, + "Destination allocation failed with an unknown exception."); + } + return allocation; + }); + + std::vector files; + HANDLE active_file = INVALID_HANDLE_VALUE; + uint64_t last_submitted_fence = 0; + bool load_succeeded = false; + // Resources referenced by in-flight I/O or GPU work must remain alive on + // every error and cancellation path. + auto cleanup = gsl::finally([&]() noexcept { + if (!load_succeeded) { + DrainActiveReads(active_file); + WaitForFenceUncancelled(last_submitted_fence); } - - hr = device_->CreateFence( - 0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(©_fence_)); - if (FAILED(hr)) { - return HResultError("ID3D12Device::CreateFence", hr); + if (allocation_future.valid()) { + try { + (void)allocation_future.get(); + } catch (...) { + } } + }); + + ORT_RETURN_IF_ERROR(PrepareFiles(ranges, cancellation, files)); - fence_event_.reset(CreateEventExW( - nullptr, nullptr, 0, EVENT_MODIFY_STATE | SYNCHRONIZE)); - if (!fence_event_) { - return Win32Error("CreateEventExW(copy fence)", GetLastError()); + uint64_t maximum_alignment = 1; + for (const auto& file : files) { + maximum_alignment = std::max(maximum_alignment, file.alignment); + } + ORT_RETURN_IF_ERROR(PrepareSlots(maximum_alignment, cancellation)); + + Batch loaded_batch; + bool allocation_ready = false; + auto ensure_allocation = [&]() -> common::Status { + if (allocation_ready) { + return common::Status::OK(); } + AllocationResult allocation = allocation_future.get(); + if (!allocation.status.IsOK()) { + return allocation.status; + } + loaded_batch = std::move(allocation.batch); + allocation_ready = true; + return common::Status::OK(); + }; - const auto upload_heap = HeapProperties(D3D12_HEAP_TYPE_UPLOAD); - const auto upload_description = - BufferDescription(config_.upload_slot_size, D3D12_RESOURCE_FLAG_NONE); - slots_.resize(config_.upload_slot_count); - for (auto& slot : slots_) { - hr = device_->CreateCommittedResource( - &upload_heap, D3D12_HEAP_FLAG_NONE, &upload_description, - D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, - IID_PPV_ARGS(&slot.resource)); - if (FAILED(hr)) { - return HResultError( - "ID3D12Device::CreateCommittedResource(upload)", hr); + for (auto& file : files) { + active_file = file.file.get(); + for (const auto& region : file.regions) { + const auto status = ReadRegion( + file, region, ranges, loaded_batch, allocation_ready, + ensure_allocation, cancellation, last_submitted_fence); + if (!status.IsOK()) { + return status; } + } + active_file = INVALID_HANDLE_VALUE; + } - hr = slot.resource->Map(0, nullptr, &slot.mapped); - if (FAILED(hr)) { - return HResultError("ID3D12Resource::Map(upload)", hr); - } + ORT_RETURN_IF_ERROR(ensure_allocation()); + ORT_RETURN_IF_ERROR( + WaitForFence(last_submitted_fence, cancellation)); + ORT_RETURN_IF_ERROR( + TransitionToCommon(loaded_batch, cancellation, last_submitted_fence)); - hr = device_->CreateCommandAllocator( - D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&slot.allocator)); - if (FAILED(hr)) { - return HResultError( - "ID3D12Device::CreateCommandAllocator", hr); - } + result = std::move(loaded_batch); + load_succeeded = true; + return common::Status::OK(); +} - hr = device_->CreateCommandList( - 0, D3D12_COMMAND_LIST_TYPE_COPY, slot.allocator.Get(), nullptr, - IID_PPV_ARGS(&slot.command_list)); - if (FAILED(hr)) { - return HResultError("ID3D12Device::CreateCommandList", hr); - } - hr = slot.command_list->Close(); - if (FAILED(hr)) { - return HResultError( - "ID3D12GraphicsCommandList::Close(initial)", hr); - } +common::Status D3D12FileBufferLoader::Impl::AllocateDestinations( + const std::vector& sizes, + Batch& batch) { + struct Placement { + bool committed = false; + size_t heap_index = 0; + uint64_t offset = 0; + D3D12_RESOURCE_DESC description{}; + }; - slot.read_event.reset(CreateEventExW( - nullptr, nullptr, CREATE_EVENT_MANUAL_RESET, - EVENT_MODIFY_STATE | SYNCHRONIZE)); - if (!slot.read_event) { - return Win32Error("CreateEventExW(read)", GetLastError()); - } + std::vector placements; + placements.reserve(sizes.size()); + std::vector heap_sizes; + + for (size_t index = 0; index < sizes.size(); ++index) { + uint64_t resource_size = 0; + if (!TryAlignUp(sizes[index], 16, resource_size)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, INVALID_ARGUMENT, + "Destination ", index, " size overflows while aligning for Dawn."); + } + const auto description = BufferDescription( + resource_size, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS); + const auto allocation_info = + device_->GetResourceAllocationInfo(0, 1, &description); + if (allocation_info.SizeInBytes == + std::numeric_limits::max()) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "GetResourceAllocationInfo failed for destination ", index, "."); + } + if (allocation_info.SizeInBytes > config_.max_destination_heap_size) { + placements.push_back({true, 0, 0, description}); + continue; } - return common::Status::OK(); - } + if (heap_sizes.empty()) { + heap_sizes.push_back(0); + } + size_t heap_index = heap_sizes.size() - 1; + uint64_t offset = 0; + if (!TryAlignUp( + heap_sizes.back(), allocation_info.Alignment, offset)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Destination placement offset overflow."); + } + if (heap_sizes.back() != 0 && + offset > + config_.max_destination_heap_size - + allocation_info.SizeInBytes) { + heap_sizes.push_back(0); + heap_index = heap_sizes.size() - 1; + offset = 0; + } - common::Status Load( - const std::vector& ranges, - Batch& result, - const CancellationToken& cancellation) { - std::lock_guard lock(load_mutex_); - result.Clear(); + placements.push_back({false, heap_index, offset, description}); + heap_sizes[heap_index] = + offset + allocation_info.SizeInBytes; + } - if (ranges.empty()) { + batch.heaps.resize(heap_sizes.size()); + for (size_t index = 0; index < heap_sizes.size(); ++index) { + uint64_t heap_size = 0; + if (!TryAlignUp( + heap_sizes[index], kBufferAlignment, heap_size) || + heap_size > config_.max_destination_heap_size) { return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "At least one file range is required."); + ONNXRUNTIME, FAIL, + "Destination heap size is invalid."); } - if (cancellation.IsCancellationRequested()) { - return CancelledStatus(); + + D3D12_HEAP_DESC heap_description{}; + heap_description.SizeInBytes = heap_size; + heap_description.Properties = + HeapProperties(D3D12_HEAP_TYPE_DEFAULT); + heap_description.Flags = D3D12_HEAP_FLAG_ALLOW_ONLY_BUFFERS; + const HRESULT hr = device_->CreateHeap( + &heap_description, IID_PPV_ARGS(&batch.heaps[index])); + if (FAILED(hr)) { + return HResultError("ID3D12Device::CreateHeap", hr); } + } - std::vector sizes; - sizes.reserve(ranges.size()); - for (size_t index = 0; index < ranges.size(); ++index) { - const auto& range = ranges[index]; - if (range.path.empty()) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "File range ", index, " has an empty path."); - } - if (range.length == 0) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "File range ", index, " has zero length."); - } - if (range.offset > - std::numeric_limits::max() - range.length) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "File range ", index, " overflows uint64_t."); - } - sizes.push_back(range.length); + batch.buffers.resize(sizes.size()); + for (size_t index = 0; index < placements.size(); ++index) { + const auto& placement = placements[index]; + HRESULT hr = S_OK; + if (placement.committed) { + const auto heap_properties = HeapProperties(D3D12_HEAP_TYPE_DEFAULT); + hr = device_->CreateCommittedResource( + &heap_properties, D3D12_HEAP_FLAG_NONE, &placement.description, + D3D12_RESOURCE_STATE_COPY_DEST, nullptr, + IID_PPV_ARGS(&batch.buffers[index].resource)); + } else { + hr = device_->CreatePlacedResource( + batch.heaps[placement.heap_index].Get(), + placement.offset, &placement.description, + D3D12_RESOURCE_STATE_COPY_DEST, nullptr, + IID_PPV_ARGS(&batch.buffers[index].resource)); + } + if (FAILED(hr)) { + return HResultError( + placement.committed + ? "ID3D12Device::CreateCommittedResource(destination)" + : "ID3D12Device::CreatePlacedResource", + hr); } + batch.buffers[index].size = sizes[index]; + } - std::future allocation_future = std::async( - std::launch::async, [this, sizes = std::move(sizes)]() noexcept { - AllocationResult allocation; - try { - allocation.status = - AllocateDestinations(sizes, allocation.batch); - } catch (const std::exception& ex) { - allocation.status = ORT_MAKE_STATUS( - ONNXRUNTIME, RUNTIME_EXCEPTION, - "Destination allocation failed: ", ex.what()); - } catch (...) { - allocation.status = ORT_MAKE_STATUS( - ONNXRUNTIME, RUNTIME_EXCEPTION, - "Destination allocation failed with an unknown exception."); - } - return allocation; - }); + return common::Status::OK(); +} - std::vector files; - HANDLE active_file = INVALID_HANDLE_VALUE; - uint64_t last_submitted_fence = 0; - bool load_succeeded = false; - auto cleanup = gsl::finally([&]() noexcept { - if (!load_succeeded) { - DrainActiveReads(active_file); - WaitForFenceUncancelled(last_submitted_fence); - } - if (allocation_future.valid()) { - try { - (void)allocation_future.get(); - } catch (...) { - } - } - }); +common::Status D3D12FileBufferLoader::Impl::PrepareFiles( + const std::vector& ranges, + const CancellationToken& cancellation, + std::vector& files) { + // Group ranges by source file so each file is opened once. Later, aligned + // overlapping ranges are merged to avoid redundant unbuffered reads. + for (size_t range_index = 0; + range_index < ranges.size(); ++range_index) { + const auto& range = ranges[range_index]; + auto file_it = std::find_if( + files.begin(), files.end(), + [&](const PreparedFile& file) { + return file.path == range.path; + }); + if (file_it == files.end()) { + PreparedFile file; + file.path = range.path; + files.push_back(std::move(file)); + file_it = std::prev(files.end()); + } + file_it->range_indices.push_back(range_index); + } - ORT_RETURN_IF_ERROR(PrepareFiles(ranges, cancellation, files)); + for (auto& file : files) { + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } - uint64_t maximum_alignment = 1; - for (const auto& file : files) { - maximum_alignment = std::max(maximum_alignment, file.alignment); + file.file.reset(CreateFileW( + file.path.c_str(), GENERIC_READ, FILE_SHARE_READ, + nullptr, OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED | + FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN, + nullptr)); + if (!file.file) { + return Win32Error( + "CreateFileW(unbuffered)", GetLastError()); } - ORT_RETURN_IF_ERROR(PrepareSlots(maximum_alignment, cancellation)); - Batch loaded_batch; - bool allocation_ready = false; - auto ensure_allocation = [&]() -> common::Status { - if (allocation_ready) { - return common::Status::OK(); - } - AllocationResult allocation = allocation_future.get(); - if (!allocation.status.IsOK()) { - return allocation.status; - } - loaded_batch = std::move(allocation.batch); - allocation_ready = true; - return common::Status::OK(); - }; - - for (auto& file : files) { - active_file = file.file.get(); - for (const auto& region : file.regions) { - const auto status = ReadRegion( - file, region, ranges, loaded_batch, allocation_ready, - ensure_allocation, cancellation, last_submitted_fence); - if (!status.IsOK()) { - return status; - } - } - active_file = INVALID_HANDLE_VALUE; + LARGE_INTEGER file_size{}; + if (!GetFileSizeEx(file.file.get(), &file_size)) { + return Win32Error("GetFileSizeEx", GetLastError()); + } + if (file_size.QuadPart < 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, "Source file has a negative size."); } + file.size = static_cast(file_size.QuadPart); - ORT_RETURN_IF_ERROR(ensure_allocation()); - ORT_RETURN_IF_ERROR( - WaitForFence(last_submitted_fence, cancellation)); - ORT_RETURN_IF_ERROR( - TransitionToCommon(loaded_batch, cancellation, last_submitted_fence)); + FILE_STORAGE_INFO storage_info{}; + if (!GetFileInformationByHandleEx( + file.file.get(), FileStorageInfo, + &storage_info, sizeof(storage_info))) { + return Win32Error( + "GetFileInformationByHandleEx(FileStorageInfo)", + GetLastError()); + } + file.alignment = std::max( + storage_info.LogicalBytesPerSector, + storage_info.PhysicalBytesPerSectorForPerformance); + if (file.alignment == 0 || + (file.alignment & (file.alignment - 1)) != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Storage alignment is not a non-zero power of two."); + } + if (config_.upload_slot_size % file.alignment != 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "upload_slot_size does not satisfy source file alignment."); + } - result = std::move(loaded_batch); - load_succeeded = true; - return common::Status::OK(); - } + std::sort( + file.range_indices.begin(), file.range_indices.end(), + [&](size_t left, size_t right) { + return ranges[left].offset < ranges[right].offset; + }); - common::Status AllocateDestinations( - const std::vector& sizes, - Batch& batch) { - struct Placement { - bool committed = false; - size_t heap_index = 0; - uint64_t offset = 0; - D3D12_RESOURCE_DESC description{}; - }; - - std::vector placements; - placements.reserve(sizes.size()); - std::vector heap_sizes; - - for (size_t index = 0; index < sizes.size(); ++index) { - uint64_t resource_size = 0; - if (!TryAlignUp(sizes[index], 16, resource_size)) { + for (size_t range_index : file.range_indices) { + const auto& range = ranges[range_index]; + if (range.offset > file.size || + range.length > file.size - range.offset) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, - "Destination ", index, " size overflows while aligning for Dawn."); - } - const auto description = BufferDescription( - resource_size, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS); - const auto allocation_info = - device_->GetResourceAllocationInfo(0, 1, &description); - if (allocation_info.SizeInBytes == - std::numeric_limits::max()) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "GetResourceAllocationInfo failed for destination ", index, "."); - } - if (allocation_info.SizeInBytes > config_.max_destination_heap_size) { - placements.push_back({true, 0, 0, description}); - continue; + "File range ", range_index, + " exceeds its source file."); } - if (heap_sizes.empty()) { - heap_sizes.push_back(0); - } - size_t heap_index = heap_sizes.size() - 1; - uint64_t offset = 0; + uint64_t aligned_end = 0; + uint64_t aligned_file_size = 0; if (!TryAlignUp( - heap_sizes.back(), allocation_info.Alignment, offset)) { + range.offset + range.length, + file.alignment, aligned_end) || + !TryAlignUp( + file.size, file.alignment, + aligned_file_size)) { return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Destination placement offset overflow."); + ONNXRUNTIME, INVALID_ARGUMENT, + "Aligned file range overflow."); } - if (heap_sizes.back() != 0 && - offset > - config_.max_destination_heap_size - - allocation_info.SizeInBytes) { - heap_sizes.push_back(0); - heap_index = heap_sizes.size() - 1; - offset = 0; + FileReadRegion region{ + AlignDown(range.offset, file.alignment), + std::min(aligned_end, aligned_file_size)}; + if (!file.regions.empty() && + region.begin <= file.regions.back().end) { + file.regions.back().end = + std::max(file.regions.back().end, region.end); + } else { + file.regions.push_back(region); } - - placements.push_back({false, heap_index, offset, description}); - heap_sizes[heap_index] = - offset + allocation_info.SizeInBytes; } + } - batch.heaps.resize(heap_sizes.size()); - for (size_t index = 0; index < heap_sizes.size(); ++index) { - uint64_t heap_size = 0; - if (!TryAlignUp( - heap_sizes[index], kBufferAlignment, heap_size) || - heap_size > config_.max_destination_heap_size) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Destination heap size is invalid."); - } + return common::Status::OK(); +} - D3D12_HEAP_DESC heap_description{}; - heap_description.SizeInBytes = heap_size; - heap_description.Properties = - HeapProperties(D3D12_HEAP_TYPE_DEFAULT); - heap_description.Flags = D3D12_HEAP_FLAG_ALLOW_ONLY_BUFFERS; - const HRESULT hr = device_->CreateHeap( - &heap_description, IID_PPV_ARGS(&batch.heaps[index])); - if (FAILED(hr)) { - return HResultError("ID3D12Device::CreateHeap", hr); - } +common::Status D3D12FileBufferLoader::Impl::PrepareSlots( + uint64_t alignment, + const CancellationToken& cancellation) { + for (auto& slot : slots_) { + ORT_RETURN_IF_ERROR( + WaitForFence(slot.fence_value, cancellation)); + if (reinterpret_cast(slot.mapped) % + alignment != + 0) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Mapped upload buffer does not satisfy source file alignment."); } + slot.overlapped = {}; + slot.file_offset = 0; + slot.requested = 0; + slot.fence_value = 0; + slot.read_active = false; + } + return common::Status::OK(); +} - batch.buffers.resize(sizes.size()); - for (size_t index = 0; index < placements.size(); ++index) { - const auto& placement = placements[index]; - HRESULT hr = S_OK; - if (placement.committed) { - const auto heap_properties = HeapProperties(D3D12_HEAP_TYPE_DEFAULT); - hr = device_->CreateCommittedResource( - &heap_properties, D3D12_HEAP_FLAG_NONE, &placement.description, - D3D12_RESOURCE_STATE_COPY_DEST, nullptr, - IID_PPV_ARGS(&batch.buffers[index].resource)); - } else { - hr = device_->CreatePlacedResource( - batch.heaps[placement.heap_index].Get(), - placement.offset, &placement.description, - D3D12_RESOURCE_STATE_COPY_DEST, nullptr, - IID_PPV_ARGS(&batch.buffers[index].resource)); - } - if (FAILED(hr)) { - return HResultError( - placement.committed - ? "ID3D12Device::CreateCommittedResource(destination)" - : "ID3D12Device::CreatePlacedResource", - hr); - } - batch.buffers[index].size = sizes[index]; +common::Status D3D12FileBufferLoader::Impl::IssueRead( + HANDLE file, + UploadSlot& slot, + uint64_t offset, + DWORD size) { + if (!ResetEvent(slot.read_event.get())) { + return Win32Error("ResetEvent(read)", GetLastError()); + } + slot.overlapped = {}; + slot.overlapped.Offset = static_cast(offset); + slot.overlapped.OffsetHigh = + static_cast(offset >> 32); + slot.overlapped.hEvent = slot.read_event.get(); + slot.file_offset = offset; + slot.requested = size; + slot.read_active = true; + if (!ReadFile( + file, slot.mapped, size, nullptr, + &slot.overlapped)) { + const DWORD error = GetLastError(); + if (error != ERROR_IO_PENDING) { + slot.read_active = false; + return Win32Error("ReadFile(unbuffered)", error); } + } + return common::Status::OK(); +} - return common::Status::OK(); +common::Status D3D12FileBufferLoader::Impl::CompleteRead( + HANDLE file, + uint64_t file_size, + UploadSlot& slot, + DWORD& bytes_read) { + if (!GetOverlappedResult( + file, &slot.overlapped, + &bytes_read, FALSE)) { + return Win32Error( + "GetOverlappedResult", GetLastError()); } + slot.read_active = false; + const uint64_t required_bytes = + std::min( + slot.requested, + file_size - slot.file_offset); + if (bytes_read < required_bytes) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Unbuffered read ended before the required file range."); + } + return common::Status::OK(); +} - common::Status PrepareFiles( - const std::vector& ranges, - const CancellationToken& cancellation, - std::vector& files) { - for (size_t range_index = 0; - range_index < ranges.size(); ++range_index) { - const auto& range = ranges[range_index]; - auto file_it = std::find_if( - files.begin(), files.end(), - [&](const PreparedFile& file) { - return file.path == range.path; - }); - if (file_it == files.end()) { - PreparedFile file; - file.path = range.path; - files.push_back(std::move(file)); - file_it = std::prev(files.end()); +template +common::Status D3D12FileBufferLoader::Impl::ReadRegion( + PreparedFile& file, + const FileReadRegion& region, + const std::vector& ranges, + Batch& batch, + bool& allocation_ready, + EnsureAllocationFn& ensure_allocation, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence) { + uint64_t next_offset = region.begin; + + // Each upload slot moves through read-active, copy-in-flight, and reusable + // states. Keep all eligible slots busy without reusing one before its + // previous GPU copy reaches the fence. + auto issue_available = [&]() -> common::Status { + for (auto& slot : slots_) { + if (next_offset >= region.end) { + break; + } + if (slot.read_active) { + continue; + } + if (slot.fence_value != 0 && + copy_fence_->GetCompletedValue() < + slot.fence_value) { + continue; } - file_it->range_indices.push_back(range_index); - } - - for (auto& file : files) { if (cancellation.IsCancellationRequested()) { return CancelledStatus(); } + const DWORD bytes = static_cast( + std::min( + config_.upload_slot_size, + region.end - next_offset)); + ORT_RETURN_IF_ERROR( + IssueRead( + file.file.get(), slot, + next_offset, bytes)); + next_offset += bytes; + } + return common::Status::OK(); + }; - file.file.reset(CreateFileW( - file.path.c_str(), GENERIC_READ, FILE_SHARE_READ, - nullptr, OPEN_EXISTING, - FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED | - FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN, - nullptr)); - if (!file.file) { - return Win32Error( - "CreateFileW(unbuffered)", GetLastError()); + ORT_RETURN_IF_ERROR(issue_available()); + while (true) { + std::array handles{}; + std::array slot_indices{}; + DWORD handle_count = 0; + for (size_t index = 0; + index < slots_.size(); ++index) { + if (slots_[index].read_active) { + handles[handle_count] = + slots_[index].read_event.get(); + slot_indices[handle_count] = index; + ++handle_count; } + } - LARGE_INTEGER file_size{}; - if (!GetFileSizeEx(file.file.get(), &file_size)) { - return Win32Error("GetFileSizeEx", GetLastError()); - } - if (file_size.QuadPart < 0) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, "Source file has a negative size."); - } - file.size = static_cast(file_size.QuadPart); - - FILE_STORAGE_INFO storage_info{}; - if (!GetFileInformationByHandleEx( - file.file.get(), FileStorageInfo, - &storage_info, sizeof(storage_info))) { - return Win32Error( - "GetFileInformationByHandleEx(FileStorageInfo)", - GetLastError()); + if (handle_count == 0) { + if (next_offset >= region.end) { + break; } - file.alignment = std::max( - storage_info.LogicalBytesPerSector, - storage_info.PhysicalBytesPerSectorForPerformance); - if (file.alignment == 0 || - (file.alignment & (file.alignment - 1)) != 0) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Storage alignment is not a non-zero power of two."); - } - if (config_.upload_slot_size % file.alignment != 0) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "upload_slot_size does not satisfy source file alignment."); - } - - std::sort( - file.range_indices.begin(), file.range_indices.end(), - [&](size_t left, size_t right) { - return ranges[left].offset < ranges[right].offset; - }); - - for (size_t range_index : file.range_indices) { - const auto& range = ranges[range_index]; - if (range.offset > file.size || - range.length > file.size - range.offset) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "File range ", range_index, - " exceeds its source file."); - } - - uint64_t aligned_end = 0; - uint64_t aligned_file_size = 0; - if (!TryAlignUp( - range.offset + range.length, - file.alignment, aligned_end) || - !TryAlignUp( - file.size, file.alignment, - aligned_file_size)) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "Aligned file range overflow."); - } - FileReadRegion region{ - AlignDown(range.offset, file.alignment), - std::min(aligned_end, aligned_file_size)}; - if (!file.regions.empty() && - region.begin <= file.regions.back().end) { - file.regions.back().end = - std::max(file.regions.back().end, region.end); - } else { - file.regions.push_back(region); + uint64_t next_fence = + std::numeric_limits::max(); + for (const auto& slot : slots_) { + if (!slot.read_active && + slot.fence_value != 0) { + next_fence = + std::min(next_fence, slot.fence_value); } } - } - - return common::Status::OK(); - } - - common::Status PrepareSlots( - uint64_t alignment, - const CancellationToken& cancellation) { - for (auto& slot : slots_) { - ORT_RETURN_IF_ERROR( - WaitForFence(slot.fence_value, cancellation)); - if (reinterpret_cast(slot.mapped) % - alignment != - 0) { + if (next_fence == + std::numeric_limits::max()) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, - "Mapped upload buffer does not satisfy source file alignment."); + "Event-driven read scheduler made no progress."); } - slot.overlapped = {}; - slot.file_offset = 0; - slot.requested = 0; - slot.fence_value = 0; - slot.read_active = false; + ORT_RETURN_IF_ERROR( + WaitForFence(next_fence, cancellation)); + ORT_RETURN_IF_ERROR(issue_available()); + continue; } - return common::Status::OK(); - } - common::Status IssueRead( - HANDLE file, - UploadSlot& slot, - uint64_t offset, - DWORD size) { - if (!ResetEvent(slot.read_event.get())) { - return Win32Error("ResetEvent(read)", GetLastError()); - } - slot.overlapped = {}; - slot.overlapped.Offset = static_cast(offset); - slot.overlapped.OffsetHigh = - static_cast(offset >> 32); - slot.overlapped.hEvent = slot.read_event.get(); - slot.file_offset = offset; - slot.requested = size; - slot.read_active = true; - if (!ReadFile( - file, slot.mapped, size, nullptr, - &slot.overlapped)) { - const DWORD error = GetLastError(); - if (error != ERROR_IO_PENDING) { - slot.read_active = false; - return Win32Error("ReadFile(unbuffered)", error); + const DWORD wait_result = WaitForMultipleObjects( + handle_count, handles.data(), FALSE, + kCancellationPollMilliseconds); + if (wait_result == WAIT_TIMEOUT) { + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); } + continue; } - return common::Status::OK(); - } - - common::Status CompleteRead( - HANDLE file, - uint64_t file_size, - UploadSlot& slot, - DWORD& bytes_read) { - if (!GetOverlappedResult( - file, &slot.overlapped, - &bytes_read, FALSE)) { + if (wait_result < WAIT_OBJECT_0 || + wait_result >= WAIT_OBJECT_0 + handle_count) { return Win32Error( - "GetOverlappedResult", GetLastError()); + "WaitForMultipleObjects(read)", + GetLastError()); } - slot.read_active = false; - const uint64_t required_bytes = - std::min( - slot.requested, - file_size - slot.file_offset); - if (bytes_read < required_bytes) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Unbuffered read ended before the required file range."); - } - return common::Status::OK(); - } - template - common::Status ReadRegion( - PreparedFile& file, - const FileReadRegion& region, - const std::vector& ranges, - Batch& batch, - bool& allocation_ready, - EnsureAllocationFn& ensure_allocation, - const CancellationToken& cancellation, - uint64_t& last_submitted_fence) { - uint64_t next_offset = region.begin; - - auto issue_available = [&]() -> common::Status { - for (auto& slot : slots_) { - if (next_offset >= region.end) { - break; - } - if (slot.read_active) { - continue; - } - if (slot.fence_value != 0 && - copy_fence_->GetCompletedValue() < - slot.fence_value) { - continue; - } - if (cancellation.IsCancellationRequested()) { - return CancelledStatus(); - } - const DWORD bytes = static_cast( - std::min( - config_.upload_slot_size, - region.end - next_offset)); - ORT_RETURN_IF_ERROR( - IssueRead( - file.file.get(), slot, - next_offset, bytes)); - next_offset += bytes; - } - return common::Status::OK(); - }; + auto& slot = + slots_[slot_indices[wait_result - WAIT_OBJECT_0]]; + DWORD bytes_read = 0; + ORT_RETURN_IF_ERROR( + CompleteRead( + file.file.get(), file.size, + slot, bytes_read)); + if (!allocation_ready) { + ORT_RETURN_IF_ERROR(ensure_allocation()); + } + ORT_RETURN_IF_ERROR( + SubmitCopies( + slot, slot.file_offset, bytes_read, + file.range_indices, ranges, batch, + last_submitted_fence)); ORT_RETURN_IF_ERROR(issue_available()); - while (true) { - std::array handles{}; - std::array slot_indices{}; - DWORD handle_count = 0; - for (size_t index = 0; - index < slots_.size(); ++index) { - if (slots_[index].read_active) { - handles[handle_count] = - slots_[index].read_event.get(); - slot_indices[handle_count] = index; - ++handle_count; - } - } - - if (handle_count == 0) { - if (next_offset >= region.end) { - break; - } - uint64_t next_fence = - std::numeric_limits::max(); - for (const auto& slot : slots_) { - if (!slot.read_active && - slot.fence_value != 0) { - next_fence = - std::min(next_fence, slot.fence_value); - } - } - if (next_fence == - std::numeric_limits::max()) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Event-driven read scheduler made no progress."); - } - ORT_RETURN_IF_ERROR( - WaitForFence(next_fence, cancellation)); - ORT_RETURN_IF_ERROR(issue_available()); - continue; - } - - const DWORD wait_result = WaitForMultipleObjects( - handle_count, handles.data(), FALSE, - kCancellationPollMilliseconds); - if (wait_result == WAIT_TIMEOUT) { - if (cancellation.IsCancellationRequested()) { - return CancelledStatus(); - } - continue; - } - if (wait_result < WAIT_OBJECT_0 || - wait_result >= WAIT_OBJECT_0 + handle_count) { - return Win32Error( - "WaitForMultipleObjects(read)", - GetLastError()); - } - - auto& slot = - slots_[slot_indices[wait_result - WAIT_OBJECT_0]]; - DWORD bytes_read = 0; - ORT_RETURN_IF_ERROR( - CompleteRead( - file.file.get(), file.size, - slot, bytes_read)); + } - if (!allocation_ready) { - ORT_RETURN_IF_ERROR(ensure_allocation()); - } - ORT_RETURN_IF_ERROR( - SubmitCopies( - slot, slot.file_offset, bytes_read, - file.range_indices, ranges, batch, - last_submitted_fence)); - ORT_RETURN_IF_ERROR(issue_available()); - } + return common::Status::OK(); +} - return common::Status::OK(); +common::Status D3D12FileBufferLoader::Impl::SubmitCopies( + UploadSlot& slot, + uint64_t chunk_begin, + uint64_t bytes, + const std::vector& range_indices, + const std::vector& ranges, + Batch& batch, + uint64_t& last_submitted_fence) { + HRESULT hr = slot.allocator->Reset(); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandAllocator::Reset", hr); + } + hr = slot.command_list->Reset( + slot.allocator.Get(), nullptr); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Reset", hr); } - common::Status SubmitCopies( - UploadSlot& slot, - uint64_t chunk_begin, - uint64_t bytes, - const std::vector& range_indices, - const std::vector& ranges, - Batch& batch, - uint64_t& last_submitted_fence) { - HRESULT hr = slot.allocator->Reset(); - if (FAILED(hr)) { - return HResultError( - "ID3D12CommandAllocator::Reset", hr); - } - hr = slot.command_list->Reset( - slot.allocator.Get(), nullptr); - if (FAILED(hr)) { - return HResultError( - "ID3D12GraphicsCommandList::Reset", hr); + const uint64_t chunk_end = chunk_begin + bytes; + for (size_t range_index : range_indices) { + const auto& range = ranges[range_index]; + const uint64_t range_end = + range.offset + range.length; + const uint64_t copy_begin = + std::max(chunk_begin, range.offset); + const uint64_t copy_end = + std::min(chunk_end, range_end); + if (copy_begin < copy_end) { + slot.command_list->CopyBufferRegion( + batch.buffers[range_index].resource.Get(), + copy_begin - range.offset, + slot.resource.Get(), + copy_begin - chunk_begin, + copy_end - copy_begin); } - const uint64_t chunk_end = chunk_begin + bytes; - for (size_t range_index : range_indices) { - const auto& range = ranges[range_index]; - const uint64_t range_end = - range.offset + range.length; - const uint64_t copy_begin = - std::max(chunk_begin, range.offset); - const uint64_t copy_end = - std::min(chunk_end, range_end); - if (copy_begin < copy_end) { + // Dawn buffer sizes are 16-byte aligned. Initialize the tail by + // repeating source bytes so no padding remains uninitialized. + const uint64_t aligned_size = + (range.length + 15) & ~uint64_t{15}; + for (uint64_t padding_offset = range.length; + padding_offset < aligned_size;) { + const uint64_t source_relative_offset = + (padding_offset - range.length) % range.length; + const uint64_t source_offset = + range.offset + source_relative_offset; + if (source_offset >= chunk_begin && + source_offset < chunk_end) { + const uint64_t copy_size = std::min( + {aligned_size - padding_offset, + range.length - source_relative_offset, + chunk_end - source_offset}); slot.command_list->CopyBufferRegion( batch.buffers[range_index].resource.Get(), - copy_begin - range.offset, + padding_offset, slot.resource.Get(), - copy_begin - chunk_begin, - copy_end - copy_begin); - } - - const uint64_t aligned_size = - (range.length + 15) & ~uint64_t{15}; - for (uint64_t padding_offset = range.length; - padding_offset < aligned_size;) { - const uint64_t source_relative_offset = - (padding_offset - range.length) % range.length; - const uint64_t source_offset = - range.offset + source_relative_offset; - if (source_offset >= chunk_begin && - source_offset < chunk_end) { - const uint64_t copy_size = std::min( - {aligned_size - padding_offset, - range.length - source_relative_offset, - chunk_end - source_offset}); - slot.command_list->CopyBufferRegion( - batch.buffers[range_index].resource.Get(), - padding_offset, - slot.resource.Get(), - source_offset - chunk_begin, - copy_size); - padding_offset += copy_size; - } else { - ++padding_offset; - } + source_offset - chunk_begin, + copy_size); + padding_offset += copy_size; + } else { + ++padding_offset; } } - - hr = slot.command_list->Close(); - if (FAILED(hr)) { - return HResultError( - "ID3D12GraphicsCommandList::Close", hr); - } - ID3D12CommandList* command_lists[] = { - slot.command_list.Get()}; - copy_queue_->ExecuteCommandLists(1, command_lists); - - slot.fence_value = ++next_fence_value_; - last_submitted_fence = slot.fence_value; - hr = copy_queue_->Signal( - copy_fence_.Get(), slot.fence_value); - if (FAILED(hr)) { - return HResultError( - "ID3D12CommandQueue::Signal", hr); - } - return common::Status::OK(); } - common::Status TransitionToCommon( - Batch& batch, - const CancellationToken& cancellation, - uint64_t& last_submitted_fence) { - if (batch.buffers.empty()) { - return common::Status::OK(); - } + hr = slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close", hr); + } + ID3D12CommandList* command_lists[] = { + slot.command_list.Get()}; + copy_queue_->ExecuteCommandLists(1, command_lists); + + slot.fence_value = ++next_fence_value_; + last_submitted_fence = slot.fence_value; + hr = copy_queue_->Signal( + copy_fence_.Get(), slot.fence_value); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandQueue::Signal", hr); + } + return common::Status::OK(); +} - auto& slot = slots_.front(); - HRESULT hr = slot.allocator->Reset(); - if (FAILED(hr)) { - return HResultError( - "ID3D12CommandAllocator::Reset(transition)", hr); - } - hr = slot.command_list->Reset( - slot.allocator.Get(), nullptr); - if (FAILED(hr)) { - return HResultError( - "ID3D12GraphicsCommandList::Reset(transition)", hr); - } +common::Status D3D12FileBufferLoader::Impl::TransitionToCommon( + Batch& batch, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence) { + if (batch.buffers.empty()) { + return common::Status::OK(); + } - std::vector barriers; - barriers.reserve(batch.buffers.size()); - for (const auto& buffer : batch.buffers) { - D3D12_RESOURCE_BARRIER barrier{}; - barrier.Type = - D3D12_RESOURCE_BARRIER_TYPE_TRANSITION; - barrier.Transition.pResource = - buffer.resource.Get(); - barrier.Transition.Subresource = - D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES; - barrier.Transition.StateBefore = - D3D12_RESOURCE_STATE_COPY_DEST; - barrier.Transition.StateAfter = - D3D12_RESOURCE_STATE_COMMON; - barriers.push_back(barrier); - } - slot.command_list->ResourceBarrier( - static_cast(barriers.size()), - barriers.data()); + auto& slot = slots_.front(); + HRESULT hr = slot.allocator->Reset(); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandAllocator::Reset(transition)", hr); + } + hr = slot.command_list->Reset( + slot.allocator.Get(), nullptr); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Reset(transition)", hr); + } - hr = slot.command_list->Close(); - if (FAILED(hr)) { - return HResultError( - "ID3D12GraphicsCommandList::Close(transition)", hr); - } - ID3D12CommandList* command_lists[] = { - slot.command_list.Get()}; - copy_queue_->ExecuteCommandLists(1, command_lists); + std::vector barriers; + barriers.reserve(batch.buffers.size()); + for (const auto& buffer : batch.buffers) { + D3D12_RESOURCE_BARRIER barrier{}; + barrier.Type = + D3D12_RESOURCE_BARRIER_TYPE_TRANSITION; + barrier.Transition.pResource = + buffer.resource.Get(); + barrier.Transition.Subresource = + D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES; + barrier.Transition.StateBefore = + D3D12_RESOURCE_STATE_COPY_DEST; + barrier.Transition.StateAfter = + D3D12_RESOURCE_STATE_COMMON; + barriers.push_back(barrier); + } + slot.command_list->ResourceBarrier( + static_cast(barriers.size()), + barriers.data()); + + hr = slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close(transition)", hr); + } + ID3D12CommandList* command_lists[] = { + slot.command_list.Get()}; + copy_queue_->ExecuteCommandLists(1, command_lists); + + last_submitted_fence = ++next_fence_value_; + hr = copy_queue_->Signal( + copy_fence_.Get(), last_submitted_fence); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandQueue::Signal(transition)", hr); + } + return WaitForFence( + last_submitted_fence, cancellation); +} - last_submitted_fence = ++next_fence_value_; - hr = copy_queue_->Signal( - copy_fence_.Get(), last_submitted_fence); - if (FAILED(hr)) { - return HResultError( - "ID3D12CommandQueue::Signal(transition)", hr); - } - return WaitForFence( - last_submitted_fence, cancellation); +common::Status D3D12FileBufferLoader::Impl::WaitForFence( + uint64_t value, + const CancellationToken& cancellation) { + if (value == 0 || + copy_fence_->GetCompletedValue() >= value) { + return common::Status::OK(); + } + const HRESULT hr = + copy_fence_->SetEventOnCompletion( + value, fence_event_.get()); + if (FAILED(hr)) { + return HResultError( + "ID3D12Fence::SetEventOnCompletion", hr); } - common::Status WaitForFence( - uint64_t value, - const CancellationToken& cancellation) { - if (value == 0 || - copy_fence_->GetCompletedValue() >= value) { + while (true) { + const DWORD wait_result = WaitForSingleObject( + fence_event_.get(), + kCancellationPollMilliseconds); + if (wait_result == WAIT_OBJECT_0) { return common::Status::OK(); } - const HRESULT hr = - copy_fence_->SetEventOnCompletion( - value, fence_event_.get()); - if (FAILED(hr)) { - return HResultError( - "ID3D12Fence::SetEventOnCompletion", hr); + if (wait_result != WAIT_TIMEOUT) { + return Win32Error( + "WaitForSingleObject(copy fence)", + GetLastError()); } - - while (true) { - const DWORD wait_result = WaitForSingleObject( - fence_event_.get(), - kCancellationPollMilliseconds); - if (wait_result == WAIT_OBJECT_0) { - return common::Status::OK(); - } - if (wait_result != WAIT_TIMEOUT) { - return Win32Error( - "WaitForSingleObject(copy fence)", - GetLastError()); - } - if (cancellation.IsCancellationRequested()) { - return CancelledStatus(); - } + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); } } +} - void WaitForFenceUncancelled(uint64_t value) noexcept { - if (value == 0 || - copy_fence_->GetCompletedValue() >= value) { - return; - } - if (SUCCEEDED(copy_fence_->SetEventOnCompletion( - value, fence_event_.get())) && - WaitForSingleObject( - fence_event_.get(), INFINITE) == - WAIT_OBJECT_0) { - return; - } - - while (true) { - const uint64_t completed = - copy_fence_->GetCompletedValue(); - if (completed >= value || - completed == std::numeric_limits::max()) { - return; - } - Sleep(1); - } +void D3D12FileBufferLoader::Impl::WaitForFenceUncancelled( + uint64_t value) noexcept { + if (value == 0 || + copy_fence_->GetCompletedValue() >= value) { + return; + } + if (SUCCEEDED(copy_fence_->SetEventOnCompletion( + value, fence_event_.get())) && + WaitForSingleObject( + fence_event_.get(), INFINITE) == + WAIT_OBJECT_0) { + return; } - void DrainActiveReads(HANDLE file) noexcept { - if (file == nullptr || - file == INVALID_HANDLE_VALUE) { + while (true) { + const uint64_t completed = + copy_fence_->GetCompletedValue(); + if (completed >= value || + completed == std::numeric_limits::max()) { return; } - for (auto& slot : slots_) { - if (!slot.read_active) { - continue; - } - (void)CancelIoEx(file, &slot.overlapped); + Sleep(1); + } +} + +void D3D12FileBufferLoader::Impl::DrainActiveReads(HANDLE file) noexcept { + if (file == nullptr || + file == INVALID_HANDLE_VALUE) { + return; + } + for (auto& slot : slots_) { + if (!slot.read_active) { + continue; } - for (auto& slot : slots_) { - if (!slot.read_active) { - continue; - } - DWORD bytes_read = 0; - (void)GetOverlappedResult( - file, &slot.overlapped, - &bytes_read, TRUE); - slot.read_active = false; + (void)CancelIoEx(file, &slot.overlapped); + } + for (auto& slot : slots_) { + if (!slot.read_active) { + continue; } + DWORD bytes_read = 0; + (void)GetOverlappedResult( + file, &slot.overlapped, + &bytes_read, TRUE); + slot.read_active = false; } - - ComPtr device_; - Config config_; - ComPtr copy_queue_; - ComPtr copy_fence_; - wil::unique_handle fence_event_; - std::vector slots_; - uint64_t next_fence_value_ = 0; - std::mutex load_mutex_; -}; +} D3D12FileBufferLoader::D3D12FileBufferLoader( std::unique_ptr impl) noexcept diff --git a/setup.py b/setup.py index 58aadc75b5010..fcde855cf20b0 100644 --- a/setup.py +++ b/setup.py @@ -468,7 +468,7 @@ def finalize_options(self): # DirectML Libs libs.extend(["DirectML.dll"]) # WebGPU/Dawn Libs - libs.extend(["dxcompiler.dll", "dxil.dll"]) + libs.extend(["dxcompiler.dll"]) # QNN V68/V73/V81 dependencies qnn_deps = [ "QnnCpu.dll", From ee9dc3bfd6be0258a4a610d1dc8dbac679d6bf9d Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:08:01 -0700 Subject: [PATCH 03/46] Refactor WebGPU context and execution provider for improved clarity and performance - Removed the unused kDeviceFreeDefaultContextId constant from webgpu_context.h. - Simplified the DeviceHasFeature method in WebGpuContext for better readability. - Updated WebGpuExecutionProvider to use the new context methods for checking device availability. - Refactored DataTransfer initialization to directly use BufferManager instead of lambda functions. - Cleaned up redundant checks and code in the WebGpuProviderFactory and related tests. - Ensured that context creation and device checks are consistent across tests. --- .../core/providers/webgpu/allocator.cc | 2 +- onnxruntime/core/providers/webgpu/allocator.h | 4 +- .../core/providers/webgpu/data_transfer.cc | 25 +- .../core/providers/webgpu/data_transfer.h | 14 +- .../core/providers/webgpu/ep/sync_stream.cc | 2 +- .../core/providers/webgpu/webgpu_context.cc | 759 +++++++++--------- .../core/providers/webgpu/webgpu_context.h | 9 +- .../webgpu/webgpu_execution_provider.cc | 13 +- .../webgpu/webgpu_provider_factory.cc | 3 - onnxruntime/core/session/inference_session.cc | 1 - .../webgpu/concurrent_context_test.cc | 6 +- .../providers/webgpu/webgpu_context_test.cc | 26 +- 12 files changed, 407 insertions(+), 457 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index 8bfed74111d97..fd210c5b5896e 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -26,7 +26,7 @@ GpuBufferAllocator::GpuBufferAllocator( buffer_manager_getter_{std::move(buffer_manager_getter)}, recording_getter_{std::move(recording_getter)}, should_submit_zero_initialize_{std::move(should_submit_zero_initialize)}, - is_read_only_allocator_{is_read_only_allocator}, + mapped_at_creation_{is_read_only_allocator && buffer_manager_getter_().SupportsUMA()}, initialize_to_zero_{!is_read_only_allocator} { } diff --git a/onnxruntime/core/providers/webgpu/allocator.h b/onnxruntime/core/providers/webgpu/allocator.h index 811daa3ac8d31..04bb783d434eb 100644 --- a/onnxruntime/core/providers/webgpu/allocator.h +++ b/onnxruntime/core/providers/webgpu/allocator.h @@ -5,7 +5,6 @@ #include #include -#include #include "core/common/narrow.h" #include "core/framework/allocator.h" @@ -58,8 +57,7 @@ class GpuBufferAllocator : public IAllocator { std::function buffer_manager_getter_; std::function recording_getter_; std::function should_submit_zero_initialize_; - bool is_read_only_allocator_; - std::optional mapped_at_creation_; + bool mapped_at_creation_; // Cached writable buffers are cleared explicitly by BufferManager::Create. Fresh buffers rely on Dawn's // "lazy_clear_resource_on_first_use" toggle, which is enabled by WebGpuContext. bool initialize_to_zero_; diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index a4681c788156a..4411c922a6fc3 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -16,27 +16,26 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data, size_t bytes) const { auto& command_state = recording_; if (bytes > 0) { - const auto& buffer_manager = buffer_manager_getter_(); if (dst_is_gpu) { if (src_is_gpu) { // copy from GPU to GPU - buffer_manager.MemCpy(command_state, - static_cast(const_cast(src_data)), - static_cast(dst_data), - bytes); + buffer_manager_.MemCpy(command_state, + static_cast(const_cast(src_data)), + static_cast(dst_data), + bytes); } else { // copy from CPU to GPU - buffer_manager.Upload(command_state, - const_cast(src_data), - static_cast(dst_data), - bytes); + buffer_manager_.Upload(command_state, + const_cast(src_data), + static_cast(dst_data), + bytes); } } else { // copy from GPU to CPU - buffer_manager.Download(command_state, - static_cast(const_cast(src_data)), - dst_data, - bytes); + buffer_manager_.Download(command_state, + static_cast(const_cast(src_data)), + dst_data, + bytes); } } diff --git a/onnxruntime/core/providers/webgpu/data_transfer.h b/onnxruntime/core/providers/webgpu/data_transfer.h index b5cc91d2a4b17..1ebc869f66d39 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.h +++ b/onnxruntime/core/providers/webgpu/data_transfer.h @@ -3,8 +3,6 @@ #pragma once -#include - #include "core/common/status.h" #include "core/framework/data_transfer.h" @@ -58,8 +56,8 @@ struct CommandRecordingState; // DataTransfer (IDataTransfer subclass) and the C API wrapper. class DataTransferImpl { public: - DataTransferImpl(std::function buffer_manager_getter, CommandRecordingState& recording) - : buffer_manager_getter_{std::move(buffer_manager_getter)}, recording_{recording} {} + DataTransferImpl(const BufferManager& buffer_manager, CommandRecordingState& recording) + : buffer_manager_{buffer_manager}, recording_{recording} {} common::Status CopyTensor(void const* src_data, bool src_is_gpu, @@ -68,16 +66,14 @@ class DataTransferImpl { size_t bytes) const; private: - std::function buffer_manager_getter_; + const BufferManager& buffer_manager_; CommandRecordingState& recording_; }; class DataTransfer : public IDataTransfer { public: - DataTransfer(std::function buffer_manager_getter, - CommandRecordingState& recording, - int context_id) - : impl_{std::move(buffer_manager_getter), recording}, context_id_{context_id} {} + DataTransfer(const BufferManager& buffer_manager, CommandRecordingState& recording, int context_id) + : impl_{buffer_manager, recording}, context_id_{context_id} {} ~DataTransfer() {}; // Device-compatibility half of CanCopy, split out because it needs no BufferManager and so can diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc index ac60452990523..f7831672483b3 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc +++ b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc @@ -83,7 +83,7 @@ CommandRecordingState& GetWebGpuStreamCommandState(const OrtSyncStream* stream) common::Status CopyTensorOnWebGpuStream(const OrtSyncStream* stream, const void* src_data, bool src_is_gpu, void* dst_data, bool dst_is_gpu, size_t bytes) { auto& ep = WebGpuSyncStream::From(stream).ep_; - DataTransferImpl transfer([&ep]() -> const BufferManager& { return ep.BufferManager(); }, ep.Recording()); + DataTransferImpl transfer(ep.BufferManager(), ep.Recording()); return transfer.CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, bytes); } diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index eafacdf6b4d47..eae82d40c8b0e 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -67,439 +67,433 @@ struct RequestAdapterOptionsLuid : wgpu::ChainedStruct { WebGpuContext::~WebGpuContext() = default; -void WebGpuContext::StartInitialize(const WebGpuContextConfig& config) { - std::call_once(init_flag_, [this, config]() { - device_free_ = config.compile_only; - Initialize(config); - }); - - if (max_num_pending_dispatches_ != config.max_num_pending_dispatches) { - LOGS_DEFAULT(WARNING) - << "WebGPU context is already initialized with " - << "maxNumPendingDispatches=" - << max_num_pending_dispatches_ - << ". Requested value " - << config.max_num_pending_dispatches - << " will be ignored."; - } - - if (config.adapter_index && - (config.adapter_index != adapter_index_ || - config.power_preference != adapter_power_preference_ || - config.backend_type != adapter_backend_type_)) { - ORT_THROW("WebGPU context is already initialized with adapterIndex=", - adapter_index_ ? std::to_string(*adapter_index_) : "automatic", - ", powerPreference=", adapter_power_preference_, - ", dawnBackendType=", adapter_backend_type_, - ". Requested selector has adapterIndex=", *config.adapter_index, - ", powerPreference=", config.power_preference, - ", dawnBackendType=", config.backend_type, - " and cannot be applied to the existing device. Use a separate process or an externally supplied custom context/device."); - } - - if (config.enable_robustness_explicitly_set) { - if (config.device != nullptr) { - LOGS_DEFAULT(WARNING) - << "WebGPU enableRobustness cannot affect an externally supplied WebGPU device. " - << "The requested value will be ignored."; - } else if (!device_free_ && enable_robustness_ != config.enable_robustness) { - LOGS_DEFAULT(WARNING) - << "WebGPU context is already initialized with enableRobustness=" << enable_robustness_ - << ". Requested value " << config.enable_robustness << " will be ignored."; - } - } -} - -#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) -ID3D12Device* WebGpuContext::WeightLoadingD3D12Device() { - return weight_loading_d3d12_device_.Get(); -} -#endif - void WebGpuContext::Initialize(const WebGpuContextConfig& config) { - max_num_pending_dispatches_ = config.max_num_pending_dispatches; - enable_robustness_ = config.enable_robustness; - adapter_index_ = config.adapter_index; - adapter_power_preference_ = config.power_preference; - adapter_backend_type_ = config.backend_type; - - ORT_ENFORCE(!config.adapter_index || config.device == nullptr, - "adapterIndex cannot be used with an externally supplied WebGPU device."); + std::call_once(init_flag_, [this, config]() { + max_num_pending_dispatches_ = config.max_num_pending_dispatches; + enable_robustness_ = config.enable_robustness; + adapter_index_ = config.adapter_index; + adapter_power_preference_ = config.power_preference; + adapter_backend_type_ = config.backend_type; + + ORT_ENFORCE(!config.adapter_index || config.device == nullptr, + "adapterIndex cannot be used with an externally supplied WebGPU device."); #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) - ORT_ENFORCE(!config.adapter_index, - "adapterIndex requires a native Dawn build with adapter enumeration support."); + ORT_ENFORCE(!config.adapter_index, + "adapterIndex requires a native Dawn build with adapter enumeration support."); #endif - weight_load_acceleration_mode_ = config.weight_load_acceleration_mode; + weight_load_acceleration_mode_ = config.weight_load_acceleration_mode; #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - requested_backend_type_ = static_cast(config.backend_type); + requested_backend_type_ = static_cast(config.backend_type); #endif - // Three easily-conflated concepts, at three layers (a pipeline, not the same flag): - // * allow_virtual_devices (env) -- selectability: surface a virtual GPU OrtEpDevice so WebGPU is - // pickable when OS enumeration finds no GPU (e.g. Win32k sandbox). - // * compile_only (session) -- intent: transform only, never finalize/run. - // * device-free / HasDevice() (ctx) -- mechanism: no Dawn device, no-op allocator. - // compile_only alone is valid (device-free even with a real GPU); a virtual device without compile_only is - // rejected at factory CreateEp -- it would try to build a real Dawn device with no hardware. - if (config.compile_only) { - device_free_ = true; - LOGS_DEFAULT(INFO) << "WebGPU EP context created device-free (compile-only session, no Dawn device)."; - return; - } + // Three easily-conflated concepts, at three layers (a pipeline, not the same flag): + // * allow_virtual_devices (env) -- selectability: surface a virtual GPU OrtEpDevice so WebGPU is + // pickable when OS enumeration finds no GPU (e.g. Win32k sandbox). + // * compile_only (session) -- intent: transform only, never finalize/run. + // * device-free / HasDevice() (ctx) -- mechanism: no Dawn device, no-op allocator. + // compile_only alone is valid (device-free even with a real GPU); a virtual device without compile_only is + // rejected at factory CreateEp -- it would try to build a real Dawn device with no hardware. + if (config.compile_only) { + LOGS_DEFAULT(INFO) << "WebGPU EP context created device-free (compile-only session, no Dawn device)."; + return; + } - if (device_ == nullptr) { - // Create wgpu::Adapter - wgpu::RequestAdapterOptions req_adapter_options = {}; - req_adapter_options.backendType = static_cast(config.backend_type); - req_adapter_options.powerPreference = static_cast(config.power_preference); + if (device_ == nullptr) { + // Create wgpu::Adapter + wgpu::RequestAdapterOptions req_adapter_options = {}; + req_adapter_options.backendType = static_cast(config.backend_type); + req_adapter_options.powerPreference = static_cast(config.power_preference); #if !defined(__wasm__) - auto enabled_adapter_toggles = GetEnabledAdapterToggles(); + auto enabled_adapter_toggles = GetEnabledAdapterToggles(); - wgpu::DawnTogglesDescriptor adapter_toggles_desc = {}; - adapter_toggles_desc.enabledToggleCount = enabled_adapter_toggles.size(); - adapter_toggles_desc.enabledToggles = enabled_adapter_toggles.data(); + wgpu::DawnTogglesDescriptor adapter_toggles_desc = {}; + adapter_toggles_desc.enabledToggleCount = enabled_adapter_toggles.size(); + adapter_toggles_desc.enabledToggles = enabled_adapter_toggles.data(); - req_adapter_options.nextInChain = &adapter_toggles_desc; + req_adapter_options.nextInChain = &adapter_toggles_desc; #endif #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - const auto* ordinary_adapter_options_chain = - req_adapter_options.nextInChain; - RequestAdapterOptionsLuid luid_options{}; - const bool preselect_weight_loading_adapter = - IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_) && - requested_backend_type_ == wgpu::BackendType::D3D12; - if (preselect_weight_loading_adapter) { - LUID selected_luid{}; - const auto selection_start = std::chrono::steady_clock::now(); - const auto selection_status = [&]() -> common::Status { - Microsoft::WRL::ComPtr dxgi_factory; - HRESULT result = - CreateDXGIFactory2(0, IID_PPV_ARGS(&dxgi_factory)); - ORT_RETURN_IF_ERROR( - FAILED(result) - ? ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Accelerated weight loading failed to create the DXGI factory " - "with HRESULT 0x", - std::hex, static_cast(result), ".") - : common::Status::OK()); - - DXGI_GPU_PREFERENCE gpu_preference = - DXGI_GPU_PREFERENCE_UNSPECIFIED; - if (req_adapter_options.powerPreference == - wgpu::PowerPreference::HighPerformance) { - gpu_preference = DXGI_GPU_PREFERENCE_HIGH_PERFORMANCE; - } else if (req_adapter_options.powerPreference == - wgpu::PowerPreference::LowPower) { - gpu_preference = DXGI_GPU_PREFERENCE_MINIMUM_POWER; - } - - for (uint32_t adapter_index = 0;; ++adapter_index) { - Microsoft::WRL::ComPtr dxgi_adapter; - result = dxgi_factory->EnumAdapterByGpuPreference( - adapter_index, gpu_preference, - IID_PPV_ARGS(&dxgi_adapter)); - if (result == DXGI_ERROR_NOT_FOUND) { - break; - } - ORT_RETURN_IF( - FAILED(result), - "Accelerated weight loading failed to enumerate DXGI adapters with " - "HRESULT 0x", - std::hex, static_cast(result), "."); - - DXGI_ADAPTER_DESC1 adapter_description{}; - result = dxgi_adapter->GetDesc1(&adapter_description); - if (FAILED(result) || - (adapter_description.Flags & DXGI_ADAPTER_FLAG_SOFTWARE) != 0) { - continue; + const auto* ordinary_adapter_options_chain = + req_adapter_options.nextInChain; + RequestAdapterOptionsLuid luid_options{}; + const bool preselect_weight_loading_adapter = + IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_) && + requested_backend_type_ == wgpu::BackendType::D3D12; + if (preselect_weight_loading_adapter) { + LUID selected_luid{}; + const auto selection_start = std::chrono::steady_clock::now(); + const auto selection_status = [&]() -> common::Status { + Microsoft::WRL::ComPtr dxgi_factory; + HRESULT result = + CreateDXGIFactory2(0, IID_PPV_ARGS(&dxgi_factory)); + ORT_RETURN_IF_ERROR( + FAILED(result) + ? ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Accelerated weight loading failed to create the DXGI factory " + "with HRESULT 0x", + std::hex, static_cast(result), ".") + : common::Status::OK()); + + DXGI_GPU_PREFERENCE gpu_preference = + DXGI_GPU_PREFERENCE_UNSPECIFIED; + if (req_adapter_options.powerPreference == + wgpu::PowerPreference::HighPerformance) { + gpu_preference = DXGI_GPU_PREFERENCE_HIGH_PERFORMANCE; + } else if (req_adapter_options.powerPreference == + wgpu::PowerPreference::LowPower) { + gpu_preference = DXGI_GPU_PREFERENCE_MINIMUM_POWER; } - Microsoft::WRL::ComPtr candidate_device; - result = D3D12CreateDevice( - dxgi_adapter.Get(), D3D_FEATURE_LEVEL_11_0, - IID_PPV_ARGS(&candidate_device)); - if (SUCCEEDED(result)) { - selected_luid = adapter_description.AdapterLuid; - weight_loading_d3d12_device_ = std::move(candidate_device); - return common::Status::OK(); - } - } - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "Accelerated weight loading did not find a hardware D3D12 adapter."); - }(); + for (uint32_t adapter_index = 0;; ++adapter_index) { + Microsoft::WRL::ComPtr dxgi_adapter; + result = dxgi_factory->EnumAdapterByGpuPreference( + adapter_index, gpu_preference, + IID_PPV_ARGS(&dxgi_adapter)); + if (result == DXGI_ERROR_NOT_FOUND) { + break; + } + ORT_RETURN_IF( + FAILED(result), + "Accelerated weight loading failed to enumerate DXGI adapters with " + "HRESULT 0x", + std::hex, static_cast(result), "."); + + DXGI_ADAPTER_DESC1 adapter_description{}; + result = dxgi_adapter->GetDesc1(&adapter_description); + if (FAILED(result) || + (adapter_description.Flags & DXGI_ADAPTER_FLAG_SOFTWARE) != 0) { + continue; + } - if (!selection_status.IsOK()) { - if (IsWeightLoadAccelerationRequired( - weight_load_acceleration_mode_)) { - ORT_THROW(selection_status.ErrorMessage()); + Microsoft::WRL::ComPtr candidate_device; + result = D3D12CreateDevice( + dxgi_adapter.Get(), D3D_FEATURE_LEVEL_11_0, + IID_PPV_ARGS(&candidate_device)); + if (SUCCEEDED(result)) { + selected_luid = adapter_description.AdapterLuid; + weight_loading_d3d12_device_ = std::move(candidate_device); + return common::Status::OK(); + } + } + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Accelerated weight loading did not find a hardware D3D12 adapter."); + }(); + + if (!selection_status.IsOK()) { + if (IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_)) { + ORT_THROW(selection_status.ErrorMessage()); + } + LOGS_DEFAULT(WARNING) + << selection_status.ErrorMessage() + << " Continuing with ordinary Dawn adapter selection without " + "accelerated weight loading."; + weight_loading_d3d12_device_.Reset(); } - LOGS_DEFAULT(WARNING) - << selection_status.ErrorMessage() - << " Continuing with ordinary Dawn adapter selection without " - "accelerated weight loading."; - weight_loading_d3d12_device_.Reset(); - } - if (weight_loading_d3d12_device_) { - luid_options.adapter_luid = selected_luid; + if (weight_loading_d3d12_device_) { + luid_options.adapter_luid = selected_luid; #if !defined(__wasm__) - luid_options.nextInChain = &adapter_toggles_desc; + luid_options.nextInChain = &adapter_toggles_desc; #endif - req_adapter_options.nextInChain = &luid_options; - - LOGS_DEFAULT(INFO) - << "WebGPU D3D12 weight-loading GPU selection: " - << std::chrono::duration( - std::chrono::steady_clock::now() - selection_start) - .count() - << " ms."; + req_adapter_options.nextInChain = &luid_options; + + LOGS_DEFAULT(INFO) + << "WebGPU D3D12 weight-loading GPU selection: " + << std::chrono::duration( + std::chrono::steady_clock::now() - selection_start) + .count() + << " ms."; + } } - } #endif - wgpu::Adapter adapter; - if (config.adapter_index) { + wgpu::Adapter adapter; + if (config.adapter_index) { #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) - const auto adapters = EnumerateBundledDawnAdapters(instance_.Get(), req_adapter_options); - ORT_ENFORCE(*config.adapter_index < adapters.size(), - "WebGPU adapterIndex ", *config.adapter_index, - " is out of range; Dawn enumerated ", adapters.size(), - " adapter(s) for the requested backend and power-preference hint."); - adapter = adapters[*config.adapter_index]; - LOGS_DEFAULT(INFO) << "WebGPU EP selected physical adapter index " << *config.adapter_index - << " of " << adapters.size() - << " adapter(s) for the requested backend and power-preference hint."; + const auto adapters = EnumerateBundledDawnAdapters(instance_.Get(), req_adapter_options); + ORT_ENFORCE(*config.adapter_index < adapters.size(), + "WebGPU adapterIndex ", *config.adapter_index, + " is out of range; Dawn enumerated ", adapters.size(), + " adapter(s) for the requested backend and power-preference hint."); + adapter = adapters[*config.adapter_index]; + LOGS_DEFAULT(INFO) << "WebGPU EP selected physical adapter index " << *config.adapter_index + << " of " << adapters.size() + << " adapter(s) for the requested backend and power-preference hint."; #endif - } else { - // Capture adapter request result without throwing inside the Dawn callback. - // Throwing C++ exceptions inside Dawn callbacks leaves Dawn's internal mutexes locked, - // which causes a self-deadlock when the WGPUInstance is later released (e.g., during - // OrtEnv teardown via EventManager::ShutDown()). - struct RequestAdapterResult { - wgpu::RequestAdapterStatus status = wgpu::RequestAdapterStatus::Error; - wgpu::Adapter adapter; - std::string message; - }; - const auto request_adapter = [&]() { - RequestAdapterResult result; - ORT_ENFORCE(wgpu::WaitStatus::Success == - instance_.WaitAny(instance_.RequestAdapter( - &req_adapter_options, - wgpu::CallbackMode::WaitAnyOnly, - [](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter, - wgpu::StringView message, - RequestAdapterResult* result) noexcept { - result->status = status; - if (status == wgpu::RequestAdapterStatus::Success) { - result->adapter = std::move(adapter); - } else { - result->message = std::string{message}; - } - }, - &result), - UINT64_MAX)); - return result; - }; - RequestAdapterResult adapter_result = request_adapter(); + } else { + // Capture adapter request result without throwing inside the Dawn callback. + // Throwing C++ exceptions inside Dawn callbacks leaves Dawn's internal mutexes locked, + // which causes a self-deadlock when the WGPUInstance is later released (e.g., during + // OrtEnv teardown via EventManager::ShutDown()). + struct RequestAdapterResult { + wgpu::RequestAdapterStatus status = wgpu::RequestAdapterStatus::Error; + wgpu::Adapter adapter; + std::string message; + }; + const auto request_adapter = [&]() { + RequestAdapterResult result; + ORT_ENFORCE(wgpu::WaitStatus::Success == + instance_.WaitAny(instance_.RequestAdapter( + &req_adapter_options, + wgpu::CallbackMode::WaitAnyOnly, + [](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter, + wgpu::StringView message, + RequestAdapterResult* result) noexcept { + result->status = status; + if (status == wgpu::RequestAdapterStatus::Success) { + result->adapter = std::move(adapter); + } else { + result->message = std::string{message}; + } + }, + &result), + UINT64_MAX)); + return result; + }; + RequestAdapterResult adapter_result = request_adapter(); #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - if (adapter_result.status != wgpu::RequestAdapterStatus::Success && - weight_loading_d3d12_device_ && - !IsWeightLoadAccelerationRequired( - weight_load_acceleration_mode_)) { - LOGS_DEFAULT(WARNING) - << "Dawn rejected the adapter selected for accelerated weight " - "loading: " - << adapter_result.message - << " Retrying ordinary Dawn adapter selection without " - "accelerated weight loading."; - weight_loading_d3d12_device_.Reset(); - req_adapter_options.nextInChain = ordinary_adapter_options_chain; - adapter_result = request_adapter(); - } + if (adapter_result.status != wgpu::RequestAdapterStatus::Success && + weight_loading_d3d12_device_ && + !IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_)) { + LOGS_DEFAULT(WARNING) + << "Dawn rejected the adapter selected for accelerated weight " + "loading: " + << adapter_result.message + << " Retrying ordinary Dawn adapter selection without " + "accelerated weight loading."; + weight_loading_d3d12_device_.Reset(); + req_adapter_options.nextInChain = ordinary_adapter_options_chain; + adapter_result = request_adapter(); + } #endif - ORT_ENFORCE(adapter_result.status == wgpu::RequestAdapterStatus::Success, - "Failed to get a WebGPU adapter: ", adapter_result.message); - adapter = std::move(adapter_result.adapter); - } - ORT_ENFORCE(adapter != nullptr, "Failed to get a WebGPU adapter."); + ORT_ENFORCE(adapter_result.status == wgpu::RequestAdapterStatus::Success, + "Failed to get a WebGPU adapter: ", adapter_result.message); + adapter = std::move(adapter_result.adapter); + } + ORT_ENFORCE(adapter != nullptr, "Failed to get a WebGPU adapter."); #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - d3d12_shared_resource_features_available_ = - adapter.HasFeature(wgpu::FeatureName::SharedBufferMemoryD3D12Resource) && - adapter.HasFeature(wgpu::FeatureName::SharedFenceDXGISharedHandle); + d3d12_shared_resource_features_available_ = + adapter.HasFeature(wgpu::FeatureName::SharedBufferMemoryD3D12Resource) && + adapter.HasFeature(wgpu::FeatureName::SharedFenceDXGISharedHandle); #endif - // Create wgpu::Device - wgpu::DeviceDescriptor device_desc = {}; + // Create wgpu::Device + wgpu::DeviceDescriptor device_desc = {}; #if !defined(__wasm__) - wgpu::DawnTogglesDescriptor device_toggles_desc = {}; - device_desc.nextInChain = &device_toggles_desc; + wgpu::DawnTogglesDescriptor device_toggles_desc = {}; + device_desc.nextInChain = &device_toggles_desc; - auto enabled_device_toggles = GetEnabledDeviceToggles(); - device_toggles_desc.enabledToggleCount = enabled_device_toggles.size(); - device_toggles_desc.enabledToggles = enabled_device_toggles.data(); + auto enabled_device_toggles = GetEnabledDeviceToggles(); + device_toggles_desc.enabledToggleCount = enabled_device_toggles.size(); + device_toggles_desc.enabledToggles = enabled_device_toggles.data(); - auto disabled_device_toggles = GetDisabledDeviceToggles(); - device_toggles_desc.disabledToggleCount = disabled_device_toggles.size(); - device_toggles_desc.disabledToggles = disabled_device_toggles.data(); + auto disabled_device_toggles = GetDisabledDeviceToggles(); + device_toggles_desc.disabledToggleCount = disabled_device_toggles.size(); + device_toggles_desc.disabledToggles = disabled_device_toggles.data(); #endif - std::vector required_features = GetAvailableRequiredFeatures(adapter); - if (!required_features.empty()) { - device_desc.requiredFeatures = required_features.data(); - device_desc.requiredFeatureCount = required_features.size(); - } - wgpu::Limits required_limits = GetRequiredLimits(adapter); - device_desc.requiredLimits = &required_limits; - - // TODO: revise temporary error handling - device_desc.SetUncapturedErrorCallback( - // Note: Don't throw from a Dawn callback. - [](const wgpu::Device& /*device*/, wgpu::ErrorType type, - wgpu::StringView message) noexcept { - if (logging::LoggingManager::HasDefaultLogger()) { - LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; - } - }); - // TODO: revise temporary device lost handling - device_desc.SetDeviceLostCallback( - wgpu::CallbackMode::AllowSpontaneous, - // Note: Don't throw from a Dawn callback. - [](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, wgpu::StringView message) noexcept { - if (logging::LoggingManager::HasDefaultLogger()) { - LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; - } - }); + std::vector required_features = GetAvailableRequiredFeatures(adapter); + if (!required_features.empty()) { + device_desc.requiredFeatures = required_features.data(); + device_desc.requiredFeatureCount = required_features.size(); + } + wgpu::Limits required_limits = GetRequiredLimits(adapter); + device_desc.requiredLimits = &required_limits; - struct RequestDeviceResult { - wgpu::RequestDeviceStatus status = wgpu::RequestDeviceStatus::Error; - wgpu::Device device; - std::string message; - }; - RequestDeviceResult device_result; - ORT_ENFORCE(wgpu::WaitStatus::Success == instance_.WaitAny(adapter.RequestDevice( - &device_desc, - wgpu::CallbackMode::WaitAnyOnly, - // Note: Don't throw from a Dawn callback. - [](wgpu::RequestDeviceStatus status, - wgpu::Device device, - wgpu::StringView message, - RequestDeviceResult* result) noexcept { - result->status = status; - if (status == wgpu::RequestDeviceStatus::Success) { - result->device = std::move(device); - } else { - result->message = std::string{message}; - } - }, - &device_result), - UINT64_MAX)); - ORT_ENFORCE(device_result.status == wgpu::RequestDeviceStatus::Success, - "Failed to get a WebGPU device: ", device_result.message); - device_ = std::move(device_result.device); - ORT_ENFORCE(device_ != nullptr, "Failed to get a WebGPU device."); - } + // TODO: revise temporary error handling + device_desc.SetUncapturedErrorCallback( + // Note: Don't throw from a Dawn callback. + [](const wgpu::Device& /*device*/, wgpu::ErrorType type, + wgpu::StringView message) noexcept { + if (logging::LoggingManager::HasDefaultLogger()) { + LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; + } + }); + // TODO: revise temporary device lost handling + device_desc.SetDeviceLostCallback( + wgpu::CallbackMode::AllowSpontaneous, + // Note: Don't throw from a Dawn callback. + [](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, wgpu::StringView message) noexcept { + if (logging::LoggingManager::HasDefaultLogger()) { + LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; + } + }); -#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - if (config.device != nullptr) { - wgpu::SupportedFeatures supported_features; - device_.GetFeatures(&supported_features); - bool has_shared_buffer_memory = false; - bool has_shared_fence = false; - for (size_t i = 0; i < supported_features.featureCount; ++i) { - has_shared_buffer_memory |= - supported_features.features[i] == - wgpu::FeatureName::SharedBufferMemoryD3D12Resource; - has_shared_fence |= - supported_features.features[i] == - wgpu::FeatureName::SharedFenceDXGISharedHandle; + struct RequestDeviceResult { + wgpu::RequestDeviceStatus status = wgpu::RequestDeviceStatus::Error; + wgpu::Device device; + std::string message; + }; + RequestDeviceResult device_result; + ORT_ENFORCE(wgpu::WaitStatus::Success == instance_.WaitAny(adapter.RequestDevice( + &device_desc, + wgpu::CallbackMode::WaitAnyOnly, + // Note: Don't throw from a Dawn callback. + [](wgpu::RequestDeviceStatus status, + wgpu::Device device, + wgpu::StringView message, + RequestDeviceResult* result) noexcept { + result->status = status; + if (status == wgpu::RequestDeviceStatus::Success) { + result->device = std::move(device); + } else { + result->message = std::string{message}; + } + }, + &device_result), + UINT64_MAX)); + ORT_ENFORCE(device_result.status == wgpu::RequestDeviceStatus::Success, + "Failed to get a WebGPU device: ", device_result.message); + device_ = std::move(device_result.device); + ORT_ENFORCE(device_ != nullptr, "Failed to get a WebGPU device."); } - d3d12_shared_resource_features_available_ = - has_shared_buffer_memory && has_shared_fence; - wgpu::AdapterInfo supplied_adapter_info; - if (device_.GetAdapterInfo(&supplied_adapter_info) == wgpu::Status::Success) { - requested_backend_type_ = supplied_adapter_info.backendType; + +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + if (config.device != nullptr) { + wgpu::SupportedFeatures supported_features; + device_.GetFeatures(&supported_features); + bool has_shared_buffer_memory = false; + bool has_shared_fence = false; + for (size_t i = 0; i < supported_features.featureCount; ++i) { + has_shared_buffer_memory |= + supported_features.features[i] == + wgpu::FeatureName::SharedBufferMemoryD3D12Resource; + has_shared_fence |= + supported_features.features[i] == + wgpu::FeatureName::SharedFenceDXGISharedHandle; + } + d3d12_shared_resource_features_available_ = + has_shared_buffer_memory && has_shared_fence; + wgpu::AdapterInfo supplied_adapter_info; + if (device_.GetAdapterInfo(&supplied_adapter_info) == wgpu::Status::Success) { + requested_backend_type_ = supplied_adapter_info.backendType; + } } - } #endif - LOGS_DEFAULT(VERBOSE) << "WebGPU EP Context is created for: Instance=" << instance_.Get() << ", Device=" << device_.Get() << "."; - - // cache device queue - device_queue_ = device_.GetQueue(); - // cache device limits - ORT_ENFORCE(Device().GetLimits(&device_limits_) == wgpu::Status::Success); - if (max_storage_buffer_binding_size_ != 0) { - device_limits_.maxStorageBufferBindingSize = - std::min(device_limits_.maxStorageBufferBindingSize, max_storage_buffer_binding_size_); - } - if (max_storage_buffers_per_shader_stage_ != 0) { - ORT_ENFORCE(max_storage_buffers_per_shader_stage_ <= device_limits_.maxStorageBuffersPerShaderStage, - "maxStorageBuffersPerShaderStage exceeds the device limit"); - device_limits_.maxStorageBuffersPerShaderStage = max_storage_buffers_per_shader_stage_; - } - // Align maxStorageBufferBindingSize down to minStorageBufferOffsetAlignment so that - // buffer segment offsets are always properly aligned for WebGPU bind group creation. - if (device_limits_.minStorageBufferOffsetAlignment > 0) { - device_limits_.maxStorageBufferBindingSize -= - (device_limits_.maxStorageBufferBindingSize % device_limits_.minStorageBufferOffsetAlignment); - } - // cache device features - wgpu::SupportedFeatures supported_features; - Device().GetFeatures(&supported_features); - for (size_t i = 0; i < supported_features.featureCount; i++) { - device_features_.insert(supported_features.features[i]); - } - // cache adapter info + LOGS_DEFAULT(VERBOSE) << "WebGPU EP Context is created for: Instance=" << instance_.Get() << ", Device=" << device_.Get() << "."; + + // cache device queue + device_queue_ = device_.GetQueue(); + // cache device limits + ORT_ENFORCE(Device().GetLimits(&device_limits_) == wgpu::Status::Success); + if (max_storage_buffer_binding_size_ != 0) { + device_limits_.maxStorageBufferBindingSize = + std::min(device_limits_.maxStorageBufferBindingSize, max_storage_buffer_binding_size_); + } + if (max_storage_buffers_per_shader_stage_ != 0) { + ORT_ENFORCE(max_storage_buffers_per_shader_stage_ <= device_limits_.maxStorageBuffersPerShaderStage, + "maxStorageBuffersPerShaderStage exceeds the device limit"); + device_limits_.maxStorageBuffersPerShaderStage = max_storage_buffers_per_shader_stage_; + } + // Align maxStorageBufferBindingSize down to minStorageBufferOffsetAlignment so that + // buffer segment offsets are always properly aligned for WebGPU bind group creation. + if (device_limits_.minStorageBufferOffsetAlignment > 0) { + device_limits_.maxStorageBufferBindingSize -= + (device_limits_.maxStorageBufferBindingSize % device_limits_.minStorageBufferOffsetAlignment); + } + // cache device features + wgpu::SupportedFeatures supported_features; + Device().GetFeatures(&supported_features); + for (size_t i = 0; i < supported_features.featureCount; i++) { + device_features_.insert(supported_features.features[i]); + } + // cache adapter info #if !defined(__wasm__) - // Native sessions share device entry points outside their per-session command recordings. - ORT_ENFORCE(DeviceHasFeature(wgpu::FeatureName::ImplicitDeviceSynchronization), - config.device != nullptr - ? "WebGPU: an externally supplied native device must enable ImplicitDeviceSynchronization " - "in DeviceDescriptor.requiredFeatures when it is created." - : "WebGPU: the internally created native device is missing the required " - "ImplicitDeviceSynchronization feature."); - if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { - adapter_info_.nextInChain = &subgroup_matrix_configs_; - } + // Native sessions share device entry points outside their per-session command recordings. + ORT_ENFORCE(DeviceHasFeature(wgpu::FeatureName::ImplicitDeviceSynchronization), + config.device != nullptr + ? "WebGPU: an externally supplied native device must enable ImplicitDeviceSynchronization " + "in DeviceDescriptor.requiredFeatures when it is created." + : "WebGPU: the internally created native device is missing the required " + "ImplicitDeviceSynchronization feature."); + if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { + adapter_info_.nextInChain = &subgroup_matrix_configs_; + } #endif - ORT_ENFORCE(Device().GetAdapterInfo(&adapter_info_) == wgpu::Status::Success); + ORT_ENFORCE(Device().GetAdapterInfo(&adapter_info_) == wgpu::Status::Success); - // create buffer manager - buffer_mgr_ = BufferManagerFactory::Create(*this, - config.buffer_cache_config.storage.mode, - config.buffer_cache_config.uniform.mode, - config.buffer_cache_config.query_resolve.mode, - config.buffer_cache_config.default_entry.mode); + // create buffer manager + buffer_mgr_ = BufferManagerFactory::Create(*this, + config.buffer_cache_config.storage.mode, + config.buffer_cache_config.uniform.mode, + config.buffer_cache_config.query_resolve.mode, + config.buffer_cache_config.default_entry.mode); - // create initializer buffer manager. - initializer_buffer_mgr_ = BufferManagerFactory::Create(*this, - BufferCacheMode::LazyRelease, - BufferCacheMode::LazyRelease, - BufferCacheMode::Disabled, - BufferCacheMode::Disabled); + // create initializer buffer manager. + initializer_buffer_mgr_ = BufferManagerFactory::Create(*this, + BufferCacheMode::LazyRelease, + BufferCacheMode::LazyRelease, + BufferCacheMode::Disabled, + BufferCacheMode::Disabled); - // create program manager - program_mgr_ = std::make_unique(*this); + // create program manager + program_mgr_ = std::make_unique(*this); - // create split-k config - split_k_config_ = std::make_unique(adapter_info_); + // create split-k config + split_k_config_ = std::make_unique(adapter_info_); - // set query type + // set query type #if !defined(__wasm__) - if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalTimestampQueryInsidePasses)) { - query_type_ = TimestampQueryType::InsidePasses; - } else + if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalTimestampQueryInsidePasses)) { + query_type_ = TimestampQueryType::InsidePasses; + } else #endif - if (DeviceHasFeature(wgpu::FeatureName::TimestampQuery)) { - query_type_ = TimestampQueryType::AtPasses; - } else { - query_type_ = TimestampQueryType::None; + if (DeviceHasFeature(wgpu::FeatureName::TimestampQuery)) { + query_type_ = TimestampQueryType::AtPasses; + } else { + query_type_ = TimestampQueryType::None; + } + }); + + if (max_num_pending_dispatches_ != config.max_num_pending_dispatches) { + LOGS_DEFAULT(WARNING) + << "WebGPU context is already initialized with " + << "maxNumPendingDispatches=" + << max_num_pending_dispatches_ + << ". Requested value " + << config.max_num_pending_dispatches + << " will be ignored."; + } + + if (config.adapter_index && + (config.adapter_index != adapter_index_ || + config.power_preference != adapter_power_preference_ || + config.backend_type != adapter_backend_type_)) { + ORT_THROW("WebGPU context is already initialized with adapterIndex=", + adapter_index_ ? std::to_string(*adapter_index_) : "automatic", + ", powerPreference=", adapter_power_preference_, + ", dawnBackendType=", adapter_backend_type_, + ". Requested selector has adapterIndex=", *config.adapter_index, + ", powerPreference=", config.power_preference, + ", dawnBackendType=", config.backend_type, + " and cannot be applied to the existing device. Use a separate process or an externally supplied custom context/device."); + } + + if (config.enable_robustness_explicitly_set) { + if (config.device != nullptr) { + LOGS_DEFAULT(WARNING) + << "WebGPU enableRobustness cannot affect an externally supplied WebGPU device. " + << "The requested value will be ignored."; + } else if (device_ != nullptr && enable_robustness_ != config.enable_robustness) { + LOGS_DEFAULT(WARNING) + << "WebGPU context is already initialized with enableRobustness=" << enable_robustness_ + << ". Requested value " << config.enable_robustness << " will be ignored."; + } } } +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +ID3D12Device* WebGpuContext::WeightLoadingD3D12Device() { + return weight_loading_d3d12_device_.Get(); +} +#endif + Status WebGpuContext::Wait(wgpu::Future f) { auto status = instance_.WaitAny(f, UINT64_MAX); if (status == wgpu::WaitStatus::Success) { @@ -1633,8 +1627,7 @@ WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& co ORT_ENFORCE(default_instance_ != nullptr, "Failed to create wgpu::Instance."); } - if (context_id == 0 || context_id == kDeviceFreeDefaultContextId) { - // The runnable and device-free default contexts share the default instance but use separate cache entries. + if (context_id == 0) { ORT_ENFORCE(instance == nullptr && device == nullptr, "WebGPU EP default contexts must not have custom WebGPU instance or device."); @@ -1673,7 +1666,7 @@ WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& co // if this was the first (and only) reference, so we don't leave a zombie context in the map // that would later deadlock during Cleanup(). ORT_TRY { - it->second.context->StartInitialize(config); + it->second.context->Initialize(config); } ORT_CATCH(...) { if (--it->second.ref_count == 0) { diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index c3128ce73927b..d46a6aca97200 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -35,7 +35,6 @@ class Tensor; namespace webgpu { -inline constexpr int kDeviceFreeDefaultContextId = -1; class WebGpuContext; class ComputeContextBase; class ComputeContext; @@ -273,9 +272,7 @@ class WebGpuContext final { const wgpu::AdapterInfo& AdapterInfo() const { return adapter_info_; } const wgpu::Limits& DeviceLimits() const { return device_limits_; } - bool DeviceHasFeature(wgpu::FeatureName feature) const { - return device_features_.contains(feature); - } + bool DeviceHasFeature(wgpu::FeatureName feature) const { return device_features_.contains(feature); } const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return subgroup_matrix_configs_; } const wgpu::CommandEncoder& GetCommandEncoder(CommandRecordingState& recording) { @@ -330,7 +327,6 @@ class WebGpuContext final { // Context-level managers are shared by sessions and synchronize their buffer caches internally. webgpu::BufferManager& BufferManager() const { return *buffer_mgr_; } - webgpu::BufferManager& InitializerBufferManager() const { return *initializer_buffer_mgr_; } inline webgpu::ValidationMode ValidationMode() const { @@ -340,7 +336,6 @@ class WebGpuContext final { // False for a device-free ("virtual device") context, which has no Dawn device and can only run graph // transformation. Used to hand out a no-op allocator instead of a real GpuBufferAllocator. inline bool HasDevice() const { return device_ != nullptr; } - inline bool IsDeviceFree() const { return device_free_; } #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) ID3D12Device* WeightLoadingD3D12Device(); bool D3D12SharedResourceFeaturesAvailable() const { @@ -433,7 +428,6 @@ class WebGpuContext final { max_storage_buffers_per_shader_stage_{max_storage_buffers_per_shader_stage} {} ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WebGpuContext); - void StartInitialize(const WebGpuContextConfig& config); void Initialize(const WebGpuContextConfig& config); wgpu::BindGroup CreateBindGroup(const std::vector& bind_buffers, @@ -472,7 +466,6 @@ class WebGpuContext final { friend class WebGpuContextFactory; std::once_flag init_flag_; - bool device_free_ = false; struct DeviceErrorState { std::mutex mutex; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index fd1a267aa1fa8..a4c5ad2d4ed11 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -74,9 +74,7 @@ class Memcpy final : public OpKernel { const auto* X = ctx->Input(0); Tensor* Y = ctx->Output(0, X->Shape()); const auto& ep = *static_cast(Info().GetExecutionProvider()); - DataTransfer transfer([&ep]() -> const BufferManager& { return ep.BufferManager(); }, - ep.Recording(), - ep.GetDeviceId()); + DataTransfer transfer(ep.BufferManager(), ep.Recording(), ep.GetDeviceId()); return transfer.CopyTensor(*X, *Y); } }; @@ -671,7 +669,7 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, } std::vector WebGpuExecutionProvider::CreatePreferredAllocators() { - const bool device_free = context_.IsDeviceFree(); + const bool device_free = !context_.HasDevice(); return { // allocator for initializers #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) @@ -799,12 +797,7 @@ std::vector> WebGpuExecutionProvider::GetCapa #endif // !defined(ORT_USE_EP_API_ADAPTERS) std::unique_ptr WebGpuExecutionProvider::GetDataTransfer() const { - return std::make_unique( - [&context = context_]() -> const webgpu::BufferManager& { - return context.BufferManager(); - }, - Recording(), - context_id_); + return std::make_unique(BufferManager(), Recording(), context_id_); } #if defined(__wasm__) diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 9c79b250f4ae0..9b501235e27a2 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -303,9 +303,6 @@ WebGpuContextConfig ParseWebGpuContextConfig(const ConfigOptions& config_options // is derived from the session config kOrtSessionOptionCompileOnly, which the Compile API sets // automatically -- same signal other EPs use (e.g. NV TensorRT RTX). Not a WebGPU-specific option. config.compile_only = config_options.GetConfigOrDefault(kOrtSessionOptionCompileOnly, "0") == "1"; - if (config.compile_only && config.context_id == 0) { - config.context_id = kDeviceFreeDefaultContextId; - } std::string max_storage_buffer_binding_size_str; if (config_options.TryGetConfigEntry(kMaxStorageBufferBindingSize, max_storage_buffer_binding_size_str)) { diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 4ce943a40bb20..9d80a191ae0da 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1212,7 +1212,6 @@ common::Status InferenceSession::LoadWithLoader(std::function const webgpu::BufferManager& { return context.BufferManager(); }, recording); + webgpu::DataTransferImpl data_transfer(context.BufferManager(), recording); std::array gpu_buffers; for (auto& buffer : gpu_buffers) { @@ -488,8 +487,7 @@ TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { std::array, kThreads> data_transfers; std::array gpu_buffers; for (int t = 0; t < kThreads; ++t) { - data_transfers[t] = std::make_unique( - [&context]() -> const webgpu::BufferManager& { return context.BufferManager(); }, recordings[t]); + data_transfers[t] = std::make_unique(context.BufferManager(), recordings[t]); wgpu::BufferDescriptor desc{}; desc.size = kElements * sizeof(float); desc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index d7c1610cf6ab5..7c1483da1eb45 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -28,7 +28,6 @@ #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/webgpu_context.h" #include "core/providers/webgpu/webgpu_execution_provider.h" -#include "core/providers/webgpu/webgpu_execution_provider.h" #include "core/providers/webgpu/webgpu_provider_factory_creator.h" #include "core/providers/webgpu/webgpu_provider_options.h" #include "core/session/onnxruntime_session_options_config_keys.h" @@ -94,7 +93,7 @@ void RunWithFreshDefaultContext(TestBody test_body, bool compile_only_parent = f } existing_provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); ASSERT_NE(existing_provider, nullptr); - existing_device = webgpu::WebGpuContextFactory::GetContext(existing_provider->GetDeviceId()).Device().Get(); + existing_device = webgpu::WebGpuContextFactory::GetContext(0).Device().Get(); } EXPECT_EXIT( @@ -104,8 +103,7 @@ void RunWithFreshDefaultContext(TestBody test_body, bool compile_only_parent = f }, testing::ExitedWithCode(EXIT_SUCCESS), ""); - EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(existing_provider->GetDeviceId()).Device().Get(), - existing_device); + EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(0).Device().Get(), existing_device); #if defined(GTEST_HAS_ABSL) && !defined(GTEST_NO_ABSL_FLAGS) flag_error.clear(); EXPECT_TRUE(death_test_style_flag->ParseFrom(previous_style, &flag_error)) << flag_error; @@ -1097,7 +1095,7 @@ TEST(WebGpuContextTest, AdapterIndexAcceptsNonNegativeInteger) { auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); ASSERT_NE(ep, nullptr); - EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()).Device().Get(), nullptr); + EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(0).Device().Get(), nullptr); }, /*compile_only_parent=*/true); #endif @@ -1137,6 +1135,7 @@ TEST(WebGpuContextTest, AdapterIndexRejectsUnsupportedBuild) { #else ConfigOptions options; ORT_THROW_IF_ERROR(options.AddConfigEntry(kAdapterIndex, "0")); + ORT_THROW_IF_ERROR(options.AddConfigEntry(kOrtSessionOptionCompileOnly, "1")); try { WebGpuProviderFactoryCreator::Create(options); @@ -1746,26 +1745,11 @@ TEST(WebGpuContextTest, CompileOnlyContextDoesNotCreateDevice) { auto ep = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); ASSERT_NE(ep, nullptr); - EXPECT_EQ(ep->GetDeviceId(), webgpu::kDeviceFreeDefaultContextId); - EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()).Device().Get(), nullptr); + EXPECT_EQ(webgpu::WebGpuContextFactory::GetContext(0).Device().Get(), nullptr); }, /*compile_only_parent=*/true); } -TEST(WebGpuContextTest, CompileOnlyAndRunnableDefaultContextsAreIsolated) { - ConfigOptions compile_only_options; - ORT_THROW_IF_ERROR(compile_only_options.AddConfigEntry(kOrtSessionOptionCompileOnly, "1")); - auto compile_only_ep = WebGpuProviderFactoryCreator::Create(compile_only_options)->CreateProvider(); - ASSERT_NE(compile_only_ep, nullptr); - - ConfigOptions runnable_options; - auto runnable_ep = WebGpuProviderFactoryCreator::Create(runnable_options)->CreateProvider(); - ASSERT_NE(runnable_ep, nullptr); - - EXPECT_TRUE(webgpu::WebGpuContextFactory::GetContext(compile_only_ep->GetDeviceId()).IsDeviceFree()); - EXPECT_FALSE(webgpu::WebGpuContextFactory::GetContext(runnable_ep->GetDeviceId()).IsDeviceFree()); -} - TEST(WebGpuContextTest, EnableRobustnessIsIndependentFromValidationMode) { #if defined(__wasm__) || defined(USE_EXTERNAL_DAWN) GTEST_SKIP() << "Dawn native toggle inspection is unavailable."; From b7c2575d1bc61169aa19d9148943c462131c4849 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 06:19:37 -0700 Subject: [PATCH 04/46] Enhance D3D12 file loading support and improve weight loading tests --- .../core/providers/webgpu/webgpu_context.cc | 16 ++- .../providers/webgpu/webgpu_context_test.cc | 135 ++++++++++++++++-- 2 files changed, 137 insertions(+), 14 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index eae82d40c8b0e..df9077f17c64d 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -15,6 +15,9 @@ #include #include "core/providers/webgpu/d3d12_external_data_loader.h" +#if !defined(USE_EXTERNAL_DAWN) +#include "dawn/native/D3D12Backend.h" +#endif #endif #if defined(__GNUC__) @@ -120,7 +123,8 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { RequestAdapterOptionsLuid luid_options{}; const bool preselect_weight_loading_adapter = IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_) && - requested_backend_type_ == wgpu::BackendType::D3D12; + requested_backend_type_ == wgpu::BackendType::D3D12 && + !config.adapter_index; if (preselect_weight_loading_adapter) { LUID selected_luid{}; const auto selection_start = std::chrono::steady_clock::now(); @@ -451,6 +455,16 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { } }); +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) && !defined(USE_EXTERNAL_DAWN) + if (IsWeightLoadAccelerationEnabled(config.weight_load_acceleration_mode) && + requested_backend_type_ == wgpu::BackendType::D3D12 && + device_ != nullptr && + weight_loading_d3d12_device_ == nullptr) { + weight_loading_d3d12_device_ = + dawn::native::d3d12::GetD3D12Device(device_.Get()); + } +#endif + if (max_num_pending_dispatches_ != config.max_num_pending_dispatches) { LOGS_DEFAULT(WARNING) << "WebGPU context is already initialized with " diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 7c1483da1eb45..59921baad7ff4 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -39,6 +40,7 @@ #include "test/util/include/temp_dir.h" #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +#include "core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h" #include "core/providers/webgpu/d3d12_external_data_loader.h" #endif @@ -1188,11 +1190,56 @@ TEST(WebGpuContextTest, RequiredWeightLoadAccelerationFailsWithoutDeviceSupport) #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) TEST(WebGpuContextTest, D3D12AcceleratedCanBeEnabledAfterInitialOffSession) { - auto off_factory = WebGpuProviderFactoryCreator::Create( - WeightLoadAccelerationOptions(kWeightLoadAcceleration_Off)); - auto off_ep = off_factory->CreateProvider(); - ASSERT_NE(off_ep, nullptr); + RunWithFreshDefaultContext([]() { + auto off_factory = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Off)); + auto off_ep = off_factory->CreateProvider(); + ASSERT_NE(off_ep, nullptr); + + auto required_factory = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions(kWeightLoadAcceleration_Required)); + auto required_ep = required_factory->CreateProvider(); + ASSERT_NE(required_ep, nullptr); + + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + const auto support_status = + webgpu::CheckD3D12AcceleratedExternalWeightsSupport(context); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + auto loader = required_ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + EXPECT_STATUS_OK(loader->BeginLoad()); + loader->AbortLoad(); + }); +} +TEST(WebGpuContextTest, D3D12FileLoaderReusesUploadSlotAcrossChunks) { + constexpr uint64_t kSlotSize = 64 * 1024; + constexpr size_t kDataSize = 3 * kSlotSize; + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_upload_slot_reuse_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + + std::vector expected(kDataSize); + for (size_t index = 0; index < expected.size(); ++index) { + expected[index] = static_cast((index * 37) & 0xff); + } + { + std::ofstream stream{data_path, std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + stream.write(reinterpret_cast(expected.data()), + static_cast(expected.size())); + ASSERT_TRUE(stream.good()); + } + + auto ep = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions( + kWeightLoadAcceleration_Required)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); auto& context = webgpu::WebGpuContextFactory::GetContext(0); const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport(context); @@ -1200,14 +1247,74 @@ TEST(WebGpuContextTest, D3D12AcceleratedCanBeEnabledAfterInitialOffSession) { GTEST_SKIP() << support_status.ErrorMessage(); } - auto required_factory = WebGpuProviderFactoryCreator::Create( - WeightLoadAccelerationOptions(kWeightLoadAcceleration_Required)); - auto required_ep = required_factory->CreateProvider(); - ASSERT_NE(required_ep, nullptr); - auto loader = required_ep->GetExternalDataLoader(); - ASSERT_NE(loader, nullptr); - EXPECT_STATUS_OK(loader->BeginLoad()); - loader->AbortLoad(); + windows::d3d12::D3D12FileBufferLoader::Config config; + config.upload_slot_count = 2; + config.upload_slot_size = kSlotSize; + std::unique_ptr loader; + ASSERT_STATUS_OK(windows::d3d12::D3D12FileBufferLoader::Create( + context.WeightLoadingD3D12Device(), loader, config)); + + windows::d3d12::D3D12FileBufferLoader::Batch batch; + ASSERT_STATUS_OK(loader->Load( + {{data_path.native(), 0, kDataSize}}, batch)); + ASSERT_EQ(batch.buffers.size(), 1u); + + auto* device = context.WeightLoadingD3D12Device(); + Microsoft::WRL::ComPtr queue; + D3D12_COMMAND_QUEUE_DESC queue_desc{}; + queue_desc.Type = D3D12_COMMAND_LIST_TYPE_COPY; + ASSERT_TRUE(SUCCEEDED(device->CreateCommandQueue( + &queue_desc, IID_PPV_ARGS(&queue)))); + + Microsoft::WRL::ComPtr allocator; + ASSERT_TRUE(SUCCEEDED(device->CreateCommandAllocator( + D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&allocator)))); + Microsoft::WRL::ComPtr command_list; + ASSERT_TRUE(SUCCEEDED(device->CreateCommandList( + 0, D3D12_COMMAND_LIST_TYPE_COPY, allocator.Get(), nullptr, + IID_PPV_ARGS(&command_list)))); + + D3D12_HEAP_PROPERTIES readback_heap{}; + readback_heap.Type = D3D12_HEAP_TYPE_READBACK; + D3D12_RESOURCE_DESC readback_desc{}; + readback_desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; + readback_desc.Width = kDataSize; + readback_desc.Height = 1; + readback_desc.DepthOrArraySize = 1; + readback_desc.MipLevels = 1; + readback_desc.SampleDesc.Count = 1; + readback_desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; + Microsoft::WRL::ComPtr readback; + ASSERT_TRUE(SUCCEEDED(device->CreateCommittedResource( + &readback_heap, D3D12_HEAP_FLAG_NONE, &readback_desc, + D3D12_RESOURCE_STATE_COPY_DEST, nullptr, + IID_PPV_ARGS(&readback)))); + + command_list->CopyBufferRegion( + readback.Get(), 0, batch.buffers[0].resource.Get(), 0, kDataSize); + ASSERT_TRUE(SUCCEEDED(command_list->Close())); + ID3D12CommandList* command_lists[]{command_list.Get()}; + queue->ExecuteCommandLists(1, command_lists); + + Microsoft::WRL::ComPtr fence; + ASSERT_TRUE(SUCCEEDED(device->CreateFence( + 0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)))); + ASSERT_TRUE(SUCCEEDED(queue->Signal(fence.Get(), 1))); + if (fence->GetCompletedValue() < 1) { + HANDLE event_handle = CreateEventW(nullptr, FALSE, FALSE, nullptr); + ASSERT_NE(event_handle, nullptr); + ASSERT_TRUE(SUCCEEDED( + fence->SetEventOnCompletion(1, event_handle))); + EXPECT_EQ(WaitForSingleObject(event_handle, INFINITE), WAIT_OBJECT_0); + CloseHandle(event_handle); + } + + void* mapped = nullptr; + const D3D12_RANGE read_range{0, kDataSize}; + ASSERT_TRUE(SUCCEEDED(readback->Map(0, &read_range, &mapped))); + EXPECT_EQ(std::memcmp(mapped, expected.data(), expected.size()), 0); + const D3D12_RANGE written_range{0, 0}; + readback->Unmap(0, &written_range); } TEST(WebGpuContextTest, D3D12AcceleratedAllocatorUsesProviderRecording) { @@ -1451,6 +1558,8 @@ void RunD3D12AcceleratedExternalInitializerSessionTest( } SessionOptions session_options; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsDisableCPUEPFallback, "1")); InferenceSession session{session_options, GetEnvironment()}; auto provider = WebGpuProviderFactoryCreator::Create( WeightLoadAccelerationOptions(mode)) @@ -1496,7 +1605,7 @@ void RunD3D12AcceleratedExternalInitializerSessionTest( TEST(WebGpuContextTest, D3D12AcceleratedLoadsExternalInitializerThroughSession) { auto probe_provider = WebGpuProviderFactoryCreator::Create( WeightLoadAccelerationOptions( - kWeightLoadAcceleration_Off)) + kWeightLoadAcceleration_Required)) ->CreateProvider(); ASSERT_NE(probe_provider, nullptr); const auto support_status = webgpu::CheckD3D12AcceleratedExternalWeightsSupport( From abf77ca4e139f11284c570b88471d472a1a838ac Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 07:06:40 -0700 Subject: [PATCH 05/46] Address D3D12 loading review and CI feedback Use the selected Dawn backend for acceleration checks, reduce loader wait and grouping overhead, and make diagnostics and Node coverage more precise. Update CUDA loader test callers for the expanded interface and account for the measured Android minimal binary growth. --- js/node/test/api/simple-api-tests.ts | 17 ++++++++++++---- .../d3d12_file_buffer_loader.cc | 20 ++++++++++--------- .../core/providers/webgpu/webgpu_context.cc | 9 +++++---- .../webgpu/webgpu_provider_factory.cc | 15 +++++++++++++- 4 files changed, 43 insertions(+), 18 deletions(-) diff --git a/js/node/test/api/simple-api-tests.ts b/js/node/test/api/simple-api-tests.ts index 0f34485d76d91..e3662dae7e96b 100644 --- a/js/node/test/api/simple-api-tests.ts +++ b/js/node/test/api/simple-api-tests.ts @@ -92,6 +92,7 @@ describe('API Tests - simple API tests', () => { const values: ReadonlyArray> = ['off', 'preferred', 'required']; const model = path.join(TEST_DATA_ROOT, 'test_types_float.onnx'); + let unavailableProviderError: string | undefined; for (const weightLoadAcceleration of values) { try { @@ -100,10 +101,18 @@ describe('API Tests - simple API tests', () => { }); await session.release(); } catch (error) { - // Builds without WebGPU/D3D12 file loading and machines without compatible - // hardware can reject the provider after parsing. The regression here - // is specifically that Node must recognize and forward the option. - assert.doesNotMatch(String(error), /WebGPU EP has an unrecognized option: 'weightLoadAcceleration'/); + const message = String(error); + assert.doesNotMatch(message, /WebGPU EP has an unrecognized option: 'weightLoadAcceleration'/); + if (weightLoadAcceleration === 'off') { + // A build without WebGPU, or a machine without a usable adapter, can + // fail before an acceleration mode matters. Later modes must not hide + // a different, unexpected failure. + unavailableProviderError = message; + } else if (weightLoadAcceleration === 'required' && unavailableProviderError === undefined) { + assert.match(message, /weightLoadAcceleration|disk-to-GPU|D3D12/); + } else { + assert.equal(message, unavailableProviderError); + } } } }); diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 1e872f4886bef..70c20c5542980 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -26,7 +27,10 @@ using Microsoft::WRL::ComPtr; namespace { constexpr uint64_t kBufferAlignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT; -constexpr DWORD kCancellationPollMilliseconds = 10; +// The cancellation token is callback-based and has no waitable handle, so a +// short timed wait keeps cancellation responsive without adding noticeable +// fence-completion tail latency. +constexpr DWORD kCancellationPollMilliseconds = 1; common::Status HResultError(const char* operation, HRESULT hr) { return ORT_MAKE_STATUS( @@ -542,21 +546,19 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( std::vector& files) { // Group ranges by source file so each file is opened once. Later, aligned // overlapping ranges are merged to avoid redundant unbuffered reads. + std::unordered_map file_indices; + file_indices.reserve(ranges.size()); for (size_t range_index = 0; range_index < ranges.size(); ++range_index) { const auto& range = ranges[range_index]; - auto file_it = std::find_if( - files.begin(), files.end(), - [&](const PreparedFile& file) { - return file.path == range.path; - }); - if (file_it == files.end()) { + const auto [file_index_it, inserted] = + file_indices.try_emplace(range.path, files.size()); + if (inserted) { PreparedFile file; file.path = range.path; files.push_back(std::move(file)); - file_it = std::prev(files.end()); } - file_it->range_indices.push_back(range_index); + files[file_index_it->second].range_indices.push_back(range_index); } for (auto& file : files) { diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index df9077f17c64d..52f13f96d2eff 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -374,10 +374,6 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { } d3d12_shared_resource_features_available_ = has_shared_buffer_memory && has_shared_fence; - wgpu::AdapterInfo supplied_adapter_info; - if (device_.GetAdapterInfo(&supplied_adapter_info) == wgpu::Status::Success) { - requested_backend_type_ = supplied_adapter_info.backendType; - } } #endif LOGS_DEFAULT(VERBOSE) << "WebGPU EP Context is created for: Instance=" << instance_.Get() << ", Device=" << device_.Get() << "."; @@ -421,6 +417,11 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { } #endif ORT_ENFORCE(Device().GetAdapterInfo(&adapter_info_) == wgpu::Status::Success); +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + // Use the backend selected by Dawn rather than the requested backend. The + // request can be Undefined when Dawn performs automatic adapter selection. + requested_backend_type_ = adapter_info_.backendType; +#endif // create buffer manager buffer_mgr_ = BufferManagerFactory::Create(*this, diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 9b501235e27a2..dbac6e8b5b26e 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -47,6 +47,19 @@ struct WebGpuProviderFactory : IExecutionProviderFactory { namespace { +constexpr const char* WeightLoadAccelerationModeName( + WeightLoadAccelerationMode mode) { + switch (mode) { + case WeightLoadAccelerationMode::Off: + return "off"; + case WeightLoadAccelerationMode::Preferred: + return "preferred"; + case WeightLoadAccelerationMode::Required: + return "required"; + } + return "unknown"; +} + WeightLoadAccelerationMode ParseWeightLoadAccelerationMode( const ConfigOptions& config_options) { std::string value; @@ -189,7 +202,7 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) LOGS_DEFAULT(VERBOSE) << "WebGPU EP preferred layout: " << int(webgpu_ep_config.data_layout); LOGS_DEFAULT(VERBOSE) << "WebGPU EP graph capture enable: " << webgpu_ep_config.enable_graph_capture; LOGS_DEFAULT(VERBOSE) << "WebGPU EP weight load acceleration mode: " - << static_cast(webgpu_ep_config.weight_load_acceleration_mode); + << WeightLoadAccelerationModeName(webgpu_ep_config.weight_load_acceleration_mode); LOGS_DEFAULT(VERBOSE) << "WebGPU EP force CPU node count: " << webgpu_ep_config.force_cpu_node_names.size(); LOGS_DEFAULT(VERBOSE) << "WebGPU EP pix capture enable: " << webgpu_ep_config.enable_pix_capture; LOGS_DEFAULT(VERBOSE) << "WebGPU EP enable int64: " << webgpu_ep_config.enable_int64; From 45e49787626937490d6bbdfcd6e90bc5c7cd55eb Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 07:31:14 -0700 Subject: [PATCH 06/46] Refine D3D12 loading lifecycle and waits Avoid exposing the accelerated allocator from device-free contexts, adapt cancellation polling during long transfers, release deferred imports explicitly, and keep per-batch timing at verbose severity. --- .../d3d12_file_buffer_loader.cc | 27 ++++++++++++++----- .../webgpu/d3d12_external_data_loader.cc | 6 +++-- .../webgpu/webgpu_execution_provider.cc | 6 +++-- .../providers/webgpu/webgpu_context_test.cc | 6 ++--- 4 files changed, 31 insertions(+), 14 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 70c20c5542980..0c2aa967f5a99 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -27,10 +27,11 @@ using Microsoft::WRL::ComPtr; namespace { constexpr uint64_t kBufferAlignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT; -// The cancellation token is callback-based and has no waitable handle, so a -// short timed wait keeps cancellation responsive without adding noticeable -// fence-completion tail latency. -constexpr DWORD kCancellationPollMilliseconds = 1; +// The cancellation token is callback-based and has no waitable handle. Start +// with a short wait to avoid adding fence-completion tail latency, then back +// off during longer transfers to avoid unnecessary CPU wakeups. +constexpr DWORD kInitialCancellationPollMilliseconds = 1; +constexpr DWORD kMaximumCancellationPollMilliseconds = 10; common::Status HResultError(const char* operation, HRESULT hr) { return ORT_MAKE_STATUS( @@ -773,6 +774,8 @@ common::Status D3D12FileBufferLoader::Impl::ReadRegion( }; ORT_RETURN_IF_ERROR(issue_available()); + DWORD cancellation_poll_milliseconds = + kInitialCancellationPollMilliseconds; while (true) { std::array handles{}; std::array slot_indices{}; @@ -814,11 +817,15 @@ common::Status D3D12FileBufferLoader::Impl::ReadRegion( const DWORD wait_result = WaitForMultipleObjects( handle_count, handles.data(), FALSE, - kCancellationPollMilliseconds); + cancellation_poll_milliseconds); if (wait_result == WAIT_TIMEOUT) { if (cancellation.IsCancellationRequested()) { return CancelledStatus(); } + cancellation_poll_milliseconds = + std::min( + cancellation_poll_milliseconds * 2, + kMaximumCancellationPollMilliseconds); continue; } if (wait_result < WAIT_OBJECT_0 || @@ -844,6 +851,8 @@ common::Status D3D12FileBufferLoader::Impl::ReadRegion( slot, slot.file_offset, bytes_read, file.range_indices, ranges, batch, last_submitted_fence)); + cancellation_poll_milliseconds = + kInitialCancellationPollMilliseconds; ORT_RETURN_IF_ERROR(issue_available()); } @@ -1013,10 +1022,12 @@ common::Status D3D12FileBufferLoader::Impl::WaitForFence( "ID3D12Fence::SetEventOnCompletion", hr); } + DWORD cancellation_poll_milliseconds = + kInitialCancellationPollMilliseconds; while (true) { const DWORD wait_result = WaitForSingleObject( fence_event_.get(), - kCancellationPollMilliseconds); + cancellation_poll_milliseconds); if (wait_result == WAIT_OBJECT_0) { return common::Status::OK(); } @@ -1028,6 +1039,10 @@ common::Status D3D12FileBufferLoader::Impl::WaitForFence( if (cancellation.IsCancellationRequested()) { return CancelledStatus(); } + cancellation_poll_milliseconds = + std::min( + cancellation_poll_milliseconds * 2, + kMaximumCancellationPollMilliseconds); } } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index a9c9c9cdb783d..de52563354e9d 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -433,7 +433,9 @@ class D3D12AcceleratedWebGpuAllocator final delete allocation; }}; recording.pending_release_callbacks.emplace_back( - [deferred = std::move(deferred)]() {}); + [deferred = std::move(deferred)]() mutable { + deferred.reset(); + }); } else { EndAccessNoThrow(*imported); } @@ -752,7 +754,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( const auto import_end = Clock::now(); batch.finalized = true; - LOGS_DEFAULT(INFO) + LOGS_DEFAULT(VERBOSE) << "WebGPU D3D12 accelerated external initializer load: " << "file-to-D3D12=" << load_metrics.load_ms << " ms, Dawn import/access=" diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index a4c5ad2d4ed11..271c8b1db64f2 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -631,7 +631,8 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, [this]() -> webgpu::CommandRecordingState& { return Recording(); }, false)} { #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) if (webgpu::IsWeightLoadAccelerationEnabled( - config.weight_load_acceleration_mode)) { + config.weight_load_acceleration_mode) && + context_.HasDevice()) { accelerated_initializer_allocator_ = CreateD3D12AcceleratedWebGpuAllocator( context_, @@ -673,7 +674,8 @@ std::vector WebGpuExecutionProvider::CreatePreferredAllocators() { return { // allocator for initializers #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - accelerated_initializer_allocator_ != nullptr + context_.HasDevice() && + accelerated_initializer_allocator_ != nullptr ? accelerated_initializer_allocator_ : #endif diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 59921baad7ff4..a19bbb49c6e39 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1173,16 +1173,14 @@ TEST(WebGpuContextTest, PreferredWeightLoadAccelerationFallsBackWithoutDeviceSup ASSERT_NE(ep, nullptr); } -TEST(WebGpuContextTest, RequiredWeightLoadAccelerationFailsWithoutDeviceSupport) { +TEST(WebGpuContextTest, RequiredWeightLoadAccelerationDoesNotExposeLoaderWithoutDevice) { auto factory = WebGpuProviderFactoryCreator::Create( CompileOnlyWeightLoadAccelerationOptions( kWeightLoadAcceleration_Required)); #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) auto ep = factory->CreateProvider(); ASSERT_NE(ep, nullptr); - auto loader = ep->GetExternalDataLoader(); - ASSERT_NE(loader, nullptr); - EXPECT_FALSE(loader->BeginLoad().IsOK()); + EXPECT_EQ(ep->GetExternalDataLoader(), nullptr); #else EXPECT_THROW(factory->CreateProvider(), OnnxRuntimeException); #endif From 556d342474386a84d1cc2645fbc3bc5e113ab2e6 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 07:54:44 -0700 Subject: [PATCH 07/46] Handle unaligned D3D12 file tails safely Keep unbuffered reads for aligned chunks while using a separate overlapped buffered handle for the final partial sector. Track each slot's active handle so cancellation and completion remain correct across both paths. --- .../d3d12_file_buffer_loader.cc | 64 +++++++++++-------- 1 file changed, 38 insertions(+), 26 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 0c2aa967f5a99..6ddb847631b51 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -100,6 +100,7 @@ struct D3D12FileBufferLoader::Impl { struct PreparedFile { std::wstring path; wil::unique_hfile file; + wil::unique_hfile buffered_tail_file; uint64_t size = 0; uint64_t alignment = 0; std::vector range_indices; @@ -113,6 +114,7 @@ struct D3D12FileBufferLoader::Impl { wil::unique_handle read_event; void* mapped = nullptr; OVERLAPPED overlapped{}; + HANDLE read_file = INVALID_HANDLE_VALUE; uint64_t file_offset = 0; DWORD requested = 0; uint64_t fence_value = 0; @@ -150,7 +152,6 @@ struct D3D12FileBufferLoader::Impl { uint64_t offset, DWORD size); common::Status CompleteRead( - HANDLE file, uint64_t file_size, UploadSlot& slot, DWORD& bytes_read); @@ -180,7 +181,7 @@ struct D3D12FileBufferLoader::Impl { uint64_t value, const CancellationToken& cancellation); void WaitForFenceUncancelled(uint64_t value) noexcept; - void DrainActiveReads(HANDLE file) noexcept; + void DrainActiveReads() noexcept; ComPtr device_; Config config_; @@ -362,14 +363,13 @@ common::Status D3D12FileBufferLoader::Impl::Load( }); std::vector files; - HANDLE active_file = INVALID_HANDLE_VALUE; uint64_t last_submitted_fence = 0; bool load_succeeded = false; // Resources referenced by in-flight I/O or GPU work must remain alive on // every error and cancellation path. auto cleanup = gsl::finally([&]() noexcept { if (!load_succeeded) { - DrainActiveReads(active_file); + DrainActiveReads(); WaitForFenceUncancelled(last_submitted_fence); } if (allocation_future.valid()) { @@ -404,7 +404,6 @@ common::Status D3D12FileBufferLoader::Impl::Load( }; for (auto& file : files) { - active_file = file.file.get(); for (const auto& region : file.regions) { const auto status = ReadRegion( file, region, ranges, loaded_batch, allocation_ready, @@ -413,7 +412,6 @@ common::Status D3D12FileBufferLoader::Impl::Load( return status; } } - active_file = INVALID_HANDLE_VALUE; } ORT_RETURN_IF_ERROR(ensure_allocation()); @@ -610,6 +608,18 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( ONNXRUNTIME, FAIL, "upload_slot_size does not satisfy source file alignment."); } + if (file.size % file.alignment != 0) { + file.buffered_tail_file.reset(CreateFileW( + file.path.c_str(), GENERIC_READ, FILE_SHARE_READ, + nullptr, OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED | + FILE_FLAG_SEQUENTIAL_SCAN, + nullptr)); + if (!file.buffered_tail_file) { + return Win32Error( + "CreateFileW(buffered tail)", GetLastError()); + } + } std::sort( file.range_indices.begin(), file.range_indices.end(), @@ -628,20 +638,16 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( } uint64_t aligned_end = 0; - uint64_t aligned_file_size = 0; if (!TryAlignUp( range.offset + range.length, - file.alignment, aligned_end) || - !TryAlignUp( - file.size, file.alignment, - aligned_file_size)) { + file.alignment, aligned_end)) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, "Aligned file range overflow."); } FileReadRegion region{ AlignDown(range.offset, file.alignment), - std::min(aligned_end, aligned_file_size)}; + std::min(aligned_end, file.size)}; if (!file.regions.empty() && region.begin <= file.regions.back().end) { file.regions.back().end = @@ -669,6 +675,7 @@ common::Status D3D12FileBufferLoader::Impl::PrepareSlots( "Mapped upload buffer does not satisfy source file alignment."); } slot.overlapped = {}; + slot.read_file = INVALID_HANDLE_VALUE; slot.file_offset = 0; slot.requested = 0; slot.fence_value = 0; @@ -692,6 +699,7 @@ common::Status D3D12FileBufferLoader::Impl::IssueRead( slot.overlapped.hEvent = slot.read_event.get(); slot.file_offset = offset; slot.requested = size; + slot.read_file = file; slot.read_active = true; if (!ReadFile( file, slot.mapped, size, nullptr, @@ -699,24 +707,25 @@ common::Status D3D12FileBufferLoader::Impl::IssueRead( const DWORD error = GetLastError(); if (error != ERROR_IO_PENDING) { slot.read_active = false; - return Win32Error("ReadFile(unbuffered)", error); + slot.read_file = INVALID_HANDLE_VALUE; + return Win32Error("ReadFile", error); } } return common::Status::OK(); } common::Status D3D12FileBufferLoader::Impl::CompleteRead( - HANDLE file, uint64_t file_size, UploadSlot& slot, DWORD& bytes_read) { if (!GetOverlappedResult( - file, &slot.overlapped, + slot.read_file, &slot.overlapped, &bytes_read, FALSE)) { return Win32Error( "GetOverlappedResult", GetLastError()); } slot.read_active = false; + slot.read_file = INVALID_HANDLE_VALUE; const uint64_t required_bytes = std::min( slot.requested, @@ -724,7 +733,7 @@ common::Status D3D12FileBufferLoader::Impl::CompleteRead( if (bytes_read < required_bytes) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, - "Unbuffered read ended before the required file range."); + "Read ended before the required file range."); } return common::Status::OK(); } @@ -764,9 +773,15 @@ common::Status D3D12FileBufferLoader::Impl::ReadRegion( std::min( config_.upload_slot_size, region.end - next_offset)); + const bool use_buffered_tail = + next_offset + bytes == file.size && + bytes % file.alignment != 0; ORT_RETURN_IF_ERROR( IssueRead( - file.file.get(), slot, + use_buffered_tail + ? file.buffered_tail_file.get() + : file.file.get(), + slot, next_offset, bytes)); next_offset += bytes; } @@ -840,8 +855,7 @@ common::Status D3D12FileBufferLoader::Impl::ReadRegion( DWORD bytes_read = 0; ORT_RETURN_IF_ERROR( CompleteRead( - file.file.get(), file.size, - slot, bytes_read)); + file.size, slot, bytes_read)); if (!allocation_ready) { ORT_RETURN_IF_ERROR(ensure_allocation()); @@ -1071,16 +1085,13 @@ void D3D12FileBufferLoader::Impl::WaitForFenceUncancelled( } } -void D3D12FileBufferLoader::Impl::DrainActiveReads(HANDLE file) noexcept { - if (file == nullptr || - file == INVALID_HANDLE_VALUE) { - return; - } +void D3D12FileBufferLoader::Impl::DrainActiveReads() noexcept { for (auto& slot : slots_) { if (!slot.read_active) { continue; } - (void)CancelIoEx(file, &slot.overlapped); + (void)CancelIoEx( + slot.read_file, &slot.overlapped); } for (auto& slot : slots_) { if (!slot.read_active) { @@ -1088,9 +1099,10 @@ void D3D12FileBufferLoader::Impl::DrainActiveReads(HANDLE file) noexcept { } DWORD bytes_read = 0; (void)GetOverlappedResult( - file, &slot.overlapped, + slot.read_file, &slot.overlapped, &bytes_read, TRUE); slot.read_active = false; + slot.read_file = INVALID_HANDLE_VALUE; } } From dcbde122c942cb532cad5c5e5a10e6d92f818113 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 08:17:19 -0700 Subject: [PATCH 08/46] Expose aligned imported WebGPU buffers Create the Dawn buffer with the initialized D3D12 resource size so normalized WebGPU copies can safely access tail padding while tensors retain their logical byte length. --- .../core/providers/webgpu/d3d12_external_data_loader.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index de52563354e9d..08100b4627019 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -720,9 +720,10 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( wgpu::BufferDescriptor buffer_descriptor{}; buffer_descriptor.label = tensor.key.name.c_str(); - buffer_descriptor.size = - static_cast( - tensor.key.length); + // Expose the initialized 16-byte-aligned resource so WebGPU's + // 4-byte-normalized copies can access tensor tail padding safely. The + // Tensor retains the logical payload length. + buffer_descriptor.size = properties.size; buffer_descriptor.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | From 3a0a6612d25ab65b3b5090c2486a87f0a82075f9 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 10:54:10 -0700 Subject: [PATCH 09/46] Make WebGPU loading tests configuration-safe Accept either no loader or an unavailable loader for compile-only contexts, and validate late-registration cleanup whether loading fails during registration or deferred session initialization. --- .../providers/webgpu/webgpu_context_test.cc | 208 +++++++++++++----- 1 file changed, 150 insertions(+), 58 deletions(-) diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index a19bbb49c6e39..50243e6c7885d 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1174,19 +1174,106 @@ TEST(WebGpuContextTest, PreferredWeightLoadAccelerationFallsBackWithoutDeviceSup } TEST(WebGpuContextTest, RequiredWeightLoadAccelerationDoesNotExposeLoaderWithoutDevice) { - auto factory = WebGpuProviderFactoryCreator::Create( - CompileOnlyWeightLoadAccelerationOptions( - kWeightLoadAcceleration_Required)); + RunWithFreshDefaultContext([]() { + auto factory = WebGpuProviderFactoryCreator::Create( + CompileOnlyWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Required)); #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - auto ep = factory->CreateProvider(); - ASSERT_NE(ep, nullptr); - EXPECT_EQ(ep->GetExternalDataLoader(), nullptr); + auto ep = factory->CreateProvider(); + ASSERT_NE(ep, nullptr); + EXPECT_EQ(ep->GetExternalDataLoader(), nullptr); #else - EXPECT_THROW(factory->CreateProvider(), OnnxRuntimeException); + EXPECT_THROW(factory->CreateProvider(), OnnxRuntimeException); #endif + }, + true); } #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +void ExpectD3D12BufferContents( + ID3D12Device* device, + const windows::d3d12::D3D12FileBufferLoader::Batch& batch, + const std::vector>& expected_buffers) { + ASSERT_EQ(batch.buffers.size(), expected_buffers.size()); + + uint64_t total_size = 0; + for (size_t index = 0; index < expected_buffers.size(); ++index) { + ASSERT_LE(expected_buffers[index].size(), batch.buffers[index].size); + total_size += expected_buffers[index].size(); + } + ASSERT_GT(total_size, 0u); + + Microsoft::WRL::ComPtr queue; + D3D12_COMMAND_QUEUE_DESC queue_desc{}; + queue_desc.Type = D3D12_COMMAND_LIST_TYPE_COPY; + ASSERT_TRUE(SUCCEEDED(device->CreateCommandQueue( + &queue_desc, IID_PPV_ARGS(&queue)))); + + Microsoft::WRL::ComPtr allocator; + ASSERT_TRUE(SUCCEEDED(device->CreateCommandAllocator( + D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&allocator)))); + Microsoft::WRL::ComPtr command_list; + ASSERT_TRUE(SUCCEEDED(device->CreateCommandList( + 0, D3D12_COMMAND_LIST_TYPE_COPY, allocator.Get(), nullptr, + IID_PPV_ARGS(&command_list)))); + + D3D12_HEAP_PROPERTIES readback_heap{}; + readback_heap.Type = D3D12_HEAP_TYPE_READBACK; + D3D12_RESOURCE_DESC readback_desc{}; + readback_desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; + readback_desc.Width = total_size; + readback_desc.Height = 1; + readback_desc.DepthOrArraySize = 1; + readback_desc.MipLevels = 1; + readback_desc.SampleDesc.Count = 1; + readback_desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; + Microsoft::WRL::ComPtr readback; + ASSERT_TRUE(SUCCEEDED(device->CreateCommittedResource( + &readback_heap, D3D12_HEAP_FLAG_NONE, &readback_desc, + D3D12_RESOURCE_STATE_COPY_DEST, nullptr, + IID_PPV_ARGS(&readback)))); + + uint64_t destination_offset = 0; + for (size_t index = 0; index < expected_buffers.size(); ++index) { + command_list->CopyBufferRegion( + readback.Get(), destination_offset, + batch.buffers[index].resource.Get(), 0, + expected_buffers[index].size()); + destination_offset += expected_buffers[index].size(); + } + ASSERT_TRUE(SUCCEEDED(command_list->Close())); + ID3D12CommandList* command_lists[]{command_list.Get()}; + queue->ExecuteCommandLists(1, command_lists); + + Microsoft::WRL::ComPtr fence; + ASSERT_TRUE(SUCCEEDED(device->CreateFence( + 0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)))); + ASSERT_TRUE(SUCCEEDED(queue->Signal(fence.Get(), 1))); + if (fence->GetCompletedValue() < 1) { + HANDLE event_handle = CreateEventW(nullptr, FALSE, FALSE, nullptr); + ASSERT_NE(event_handle, nullptr); + ASSERT_TRUE(SUCCEEDED( + fence->SetEventOnCompletion(1, event_handle))); + EXPECT_EQ(WaitForSingleObject(event_handle, INFINITE), WAIT_OBJECT_0); + CloseHandle(event_handle); + } + + void* mapped = nullptr; + const D3D12_RANGE read_range{0, static_cast(total_size)}; + ASSERT_TRUE(SUCCEEDED(readback->Map(0, &read_range, &mapped))); + size_t source_offset = 0; + for (const auto& expected : expected_buffers) { + EXPECT_EQ( + std::memcmp( + static_cast(mapped) + source_offset, + expected.data(), expected.size()), + 0); + source_offset += expected.size(); + } + const D3D12_RANGE written_range{0, 0}; + readback->Unmap(0, &written_range); +} + TEST(WebGpuContextTest, D3D12AcceleratedCanBeEnabledAfterInitialOffSession) { RunWithFreshDefaultContext([]() { auto off_factory = WebGpuProviderFactoryCreator::Create( @@ -1258,61 +1345,61 @@ TEST(WebGpuContextTest, D3D12FileLoaderReusesUploadSlotAcrossChunks) { ASSERT_EQ(batch.buffers.size(), 1u); auto* device = context.WeightLoadingD3D12Device(); - Microsoft::WRL::ComPtr queue; - D3D12_COMMAND_QUEUE_DESC queue_desc{}; - queue_desc.Type = D3D12_COMMAND_LIST_TYPE_COPY; - ASSERT_TRUE(SUCCEEDED(device->CreateCommandQueue( - &queue_desc, IID_PPV_ARGS(&queue)))); - - Microsoft::WRL::ComPtr allocator; - ASSERT_TRUE(SUCCEEDED(device->CreateCommandAllocator( - D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&allocator)))); - Microsoft::WRL::ComPtr command_list; - ASSERT_TRUE(SUCCEEDED(device->CreateCommandList( - 0, D3D12_COMMAND_LIST_TYPE_COPY, allocator.Get(), nullptr, - IID_PPV_ARGS(&command_list)))); + ExpectD3D12BufferContents(device, batch, {expected}); +} - D3D12_HEAP_PROPERTIES readback_heap{}; - readback_heap.Type = D3D12_HEAP_TYPE_READBACK; - D3D12_RESOURCE_DESC readback_desc{}; - readback_desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - readback_desc.Width = kDataSize; - readback_desc.Height = 1; - readback_desc.DepthOrArraySize = 1; - readback_desc.MipLevels = 1; - readback_desc.SampleDesc.Count = 1; - readback_desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - Microsoft::WRL::ComPtr readback; - ASSERT_TRUE(SUCCEEDED(device->CreateCommittedResource( - &readback_heap, D3D12_HEAP_FLAG_NONE, &readback_desc, - D3D12_RESOURCE_STATE_COPY_DEST, nullptr, - IID_PPV_ARGS(&readback)))); +TEST(WebGpuContextTest, D3D12FileLoaderUsesMultipleHeapsAndCommittedResources) { + constexpr uint64_t kHeapSize = 64 * 1024; + const std::vector> expected_buffers{ + std::vector(1, 0x19), + std::vector(17, 0x5a), + std::vector(kHeapSize + 1, 0xc3), + }; + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_destination_allocation_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); - command_list->CopyBufferRegion( - readback.Get(), 0, batch.buffers[0].resource.Get(), 0, kDataSize); - ASSERT_TRUE(SUCCEEDED(command_list->Close())); - ID3D12CommandList* command_lists[]{command_list.Get()}; - queue->ExecuteCommandLists(1, command_lists); + std::vector ranges; + uint64_t file_offset = 0; + { + std::ofstream stream{data_path, std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + for (const auto& expected : expected_buffers) { + ranges.push_back( + {data_path.native(), file_offset, expected.size()}); + stream.write( + reinterpret_cast(expected.data()), + static_cast(expected.size())); + file_offset += expected.size(); + } + ASSERT_TRUE(stream.good()); + } - Microsoft::WRL::ComPtr fence; - ASSERT_TRUE(SUCCEEDED(device->CreateFence( - 0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)))); - ASSERT_TRUE(SUCCEEDED(queue->Signal(fence.Get(), 1))); - if (fence->GetCompletedValue() < 1) { - HANDLE event_handle = CreateEventW(nullptr, FALSE, FALSE, nullptr); - ASSERT_NE(event_handle, nullptr); - ASSERT_TRUE(SUCCEEDED( - fence->SetEventOnCompletion(1, event_handle))); - EXPECT_EQ(WaitForSingleObject(event_handle, INFINITE), WAIT_OBJECT_0); - CloseHandle(event_handle); + auto ep = WebGpuProviderFactoryCreator::Create( + WeightLoadAccelerationOptions( + kWeightLoadAcceleration_Required)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + const auto support_status = + webgpu::CheckD3D12AcceleratedExternalWeightsSupport(context); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); } - void* mapped = nullptr; - const D3D12_RANGE read_range{0, kDataSize}; - ASSERT_TRUE(SUCCEEDED(readback->Map(0, &read_range, &mapped))); - EXPECT_EQ(std::memcmp(mapped, expected.data(), expected.size()), 0); - const D3D12_RANGE written_range{0, 0}; - readback->Unmap(0, &written_range); + windows::d3d12::D3D12FileBufferLoader::Config config; + config.max_destination_heap_size = kHeapSize; + std::unique_ptr loader; + ASSERT_STATUS_OK(windows::d3d12::D3D12FileBufferLoader::Create( + context.WeightLoadingD3D12Device(), loader, config)); + + windows::d3d12::D3D12FileBufferLoader::Batch batch; + ASSERT_STATUS_OK(loader->Load(ranges, batch)); + ASSERT_EQ(batch.buffers.size(), expected_buffers.size()); + EXPECT_EQ(batch.heaps.size(), 2u); + ExpectD3D12BufferContents( + context.WeightLoadingD3D12Device(), batch, expected_buffers); } TEST(WebGpuContextTest, D3D12AcceleratedAllocatorUsesProviderRecording) { @@ -1569,8 +1656,13 @@ void RunD3D12AcceleratedExternalInitializerSessionTest( const auto memory_info = allocators.front()->Info(); ASSERT_STATUS_OK(session.Load(model_path.native())); ASSERT_TRUE(std::filesystem::remove(data_path)); - const auto status = + auto status = session.RegisterExecutionProvider(std::move(provider)); + if (status.IsOK()) { + // Some provider configurations defer external initializer loading until + // initialization instead of attempting it during late registration. + status = session.Initialize(); + } EXPECT_FALSE(status.IsOK()); EXPECT_EQ(session.GetExternalDataLoaderManager().GetExternalDataLoader( memory_info), From 49b612b4f99bbac326cccecaff5de2793bca8f19 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 11:19:22 -0700 Subject: [PATCH 10/46] Polish D3D12 loading diagnostics --- .../windows/d3d12_file_loader/d3d12_file_buffer_loader.cc | 4 +++- .../core/providers/webgpu/d3d12_external_data_loader.cc | 2 +- .../core/providers/webgpu/d3d12_external_data_loader.h | 2 +- onnxruntime/core/providers/webgpu/webgpu_context.cc | 2 +- onnxruntime/core/providers/webgpu/webgpu_context.h | 2 +- 5 files changed, 7 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 6ddb847631b51..f91dfb0c15199 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -35,7 +36,8 @@ constexpr DWORD kMaximumCancellationPollMilliseconds = 10; common::Status HResultError(const char* operation, HRESULT hr) { return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, operation, " failed with HRESULT ", + ONNXRUNTIME, FAIL, operation, " failed with HRESULT 0x", + std::setfill('0'), std::uppercase, std::hex, std::setw(8), static_cast(hr)); } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 08100b4627019..48b58c2953d04 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -295,7 +295,7 @@ void EndAccessNoThrow( } // namespace common::Status CheckD3D12AcceleratedExternalWeightsSupport( - WebGpuContext& context) { + const WebGpuContext& context) { ORT_RETURN_IF_NOT( context.RequestedBackendType() == wgpu::BackendType::D3D12, diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index 993276cdde9ab..7f10d63193f24 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -21,7 +21,7 @@ namespace webgpu { class WebGpuContext; struct CommandRecordingState; -common::Status CheckD3D12AcceleratedExternalWeightsSupport(WebGpuContext& context); +common::Status CheckD3D12AcceleratedExternalWeightsSupport(const WebGpuContext& context); common::Status ResolveWeightLoadAccelerationMode( WeightLoadAccelerationMode mode, diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 52f13f96d2eff..db7dc39b4b12d 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -504,7 +504,7 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { } #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) -ID3D12Device* WebGpuContext::WeightLoadingD3D12Device() { +ID3D12Device* WebGpuContext::WeightLoadingD3D12Device() const { return weight_loading_d3d12_device_.Get(); } #endif diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index d46a6aca97200..da323f96752fa 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -337,7 +337,7 @@ class WebGpuContext final { // transformation. Used to hand out a no-op allocator instead of a real GpuBufferAllocator. inline bool HasDevice() const { return device_ != nullptr; } #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - ID3D12Device* WeightLoadingD3D12Device(); + ID3D12Device* WeightLoadingD3D12Device() const; bool D3D12SharedResourceFeaturesAvailable() const { return d3d12_shared_resource_features_available_; } From 7fbdf83d24cf1289de3b167156e9c2d3e99fe449 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 11:23:08 -0700 Subject: [PATCH 11/46] Fix acceleration portability test Allow required WebGPU acceleration to report its own device-unavailable error and restore the Android binary threshold from main. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- js/node/test/api/simple-api-tests.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/js/node/test/api/simple-api-tests.ts b/js/node/test/api/simple-api-tests.ts index e3662dae7e96b..82a776e50b3e2 100644 --- a/js/node/test/api/simple-api-tests.ts +++ b/js/node/test/api/simple-api-tests.ts @@ -108,8 +108,10 @@ describe('API Tests - simple API tests', () => { // fail before an acceleration mode matters. Later modes must not hide // a different, unexpected failure. unavailableProviderError = message; - } else if (weightLoadAcceleration === 'required' && unavailableProviderError === undefined) { - assert.match(message, /weightLoadAcceleration|disk-to-GPU|D3D12/); + } else if (weightLoadAcceleration === 'required') { + if (message !== unavailableProviderError) { + assert.match(message, /weightLoadAcceleration|disk-to-GPU|DXGI|D3D12/); + } } else { assert.equal(message, unavailableProviderError); } From 9166268ae5e526b8cd846a12e03c25b60a8962bb Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 11:43:30 -0700 Subject: [PATCH 12/46] Scope accelerated loading lifecycle to D3D12 --- onnxruntime/core/framework/session_state_utils.cc | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 92f2b432b447a..db74b1d73c746 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -104,7 +104,9 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); if (external_data_loader) { if (external_data_loader->CreatesTensorForDevice(device)) { - ORT_RETURN_IF(memory_buffer != nullptr || alloc == nullptr, + ORT_RETURN_IF(memory_buffer != nullptr, + "An external data loader that creates tensors cannot use a preallocated buffer."); + ORT_RETURN_IF(alloc == nullptr, "An external data loader that creates tensors requires a device allocator."); tensor = Tensor{type, tensor_shape, nullptr, alloc}; } else { @@ -395,6 +397,7 @@ common::Status SaveInitializedTensors( session_options.config_options.GetConfigOrDefault( kOrtSessionOptionsUseDeviceAllocatorForInitializers, "0") == "1"; +#if defined(ENABLE_D3D12_FILE_LOADING) ORT_RETURN_IF_ERROR(external_data_loader_mgr.BeginLoad()); bool external_data_load_finalized = false; auto abort_external_data_load = gsl::finally([&]() { @@ -425,6 +428,7 @@ common::Status SaveInitializedTensors( ORT_RETURN_IF_ERROR(external_data_loader_mgr.FinalizeLoad( [&session_options]() { return session_options.IsLoadCancellationFlagSet(); })); +#endif // 3. create weight tensors based on weights buffer for (const auto& entry : id_to_initialized_tensor) { @@ -513,7 +517,13 @@ common::Status SaveInitializedTensors( #endif } +#if defined(ENABLE_D3D12_FILE_LOADING) + // AbortLoad is also the idempotent end-of-batch cleanup hook. All claimed + // allocations have transferred ownership, so this only releases metadata + // and any unclaimed resources left by a successful batch. + external_data_loader_mgr.AbortLoad(); external_data_load_finalized = true; +#endif LOGS(logger, INFO) << "Done saving initialized tensors"; return common::Status::OK(); } From fbd0819408adb868e6b4e88f5f1e6ddd0171f0ad Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:14:07 -0700 Subject: [PATCH 13/46] Gate batched loading API with D3D12 support --- onnxruntime/core/framework/external_data_loader.h | 1 + onnxruntime/core/framework/external_data_loader_manager.cc | 2 ++ onnxruntime/core/framework/external_data_loader_manager.h | 2 ++ onnxruntime/core/framework/session_state_utils.cc | 5 ++++- onnxruntime/core/framework/tensorprotoutils.h | 2 ++ .../core/providers/webgpu/d3d12_external_data_loader.cc | 2 +- onnxruntime/test/framework/external_data_loader_test.cc | 2 ++ 7 files changed, 14 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index 535de3f9a29dc..b935a028a3ce6 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -95,6 +95,7 @@ class IExternalDataLoader { // PrepareTensor(), FinalizeLoad(), or partial tensor consumption, and may be called // more than once. Implementations must therefore be noexcept and idempotent. virtual void AbortLoad() const noexcept {} +#endif // A tensor-creating loader replaces tensor with one backed by memory owned through allocator. virtual common::Status LoadTensor([[maybe_unused]] const Env& env, diff --git a/onnxruntime/core/framework/external_data_loader_manager.cc b/onnxruntime/core/framework/external_data_loader_manager.cc index 5d1a3fa1c573c..86746b4a2126c 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.cc +++ b/onnxruntime/core/framework/external_data_loader_manager.cc @@ -36,6 +36,7 @@ const IExternalDataLoader* ExternalDataLoaderManager::GetExternalDataLoader( return nullptr; } +#if defined(ENABLE_D3D12_FILE_LOADING) const IExternalDataLoader* ExternalDataLoaderManager::GetTensorCreator( const OrtDevice& target_device, int32_t tensor_data_type) const { for (const auto& external_data_loader : external_data_loaders_) { @@ -77,5 +78,6 @@ void ExternalDataLoaderManager::AbortLoad() const noexcept { external_data_loader->AbortLoad(); } } +#endif } // namespace onnxruntime diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index 5fcd6d525ebde..056d74370e9ca 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.h +++ b/onnxruntime/core/framework/external_data_loader_manager.h @@ -24,11 +24,13 @@ class ExternalDataLoaderManager { const IExternalDataLoader* GetExternalDataLoader( const OrtMemoryInfo& target_memory_info, int32_t tensor_data_type) const; +#if defined(ENABLE_D3D12_FILE_LOADING) const IExternalDataLoader* GetTensorCreator( const OrtDevice& target_device, int32_t tensor_data_type) const; common::Status BeginLoad() const; common::Status FinalizeLoad(const std::function& is_canceled) const; void AbortLoad() const noexcept; +#endif private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExternalDataLoaderManager); diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index db74b1d73c746..f473182a77d1a 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -103,13 +103,16 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st auto external_data_loader = external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); if (external_data_loader) { +#if defined(ENABLE_D3D12_FILE_LOADING) if (external_data_loader->CreatesTensorForDevice(device)) { ORT_RETURN_IF(memory_buffer != nullptr, "An external data loader that creates tensors cannot use a preallocated buffer."); ORT_RETURN_IF(alloc == nullptr, "An external data loader that creates tensors requires a device allocator."); tensor = Tensor{type, tensor_shape, nullptr, alloc}; - } else { + } else +#endif + { // if custom external data loader is used, always allocate memory on device ORT_RETURN_IF_ERROR( AllocateTensor(memory_buffer, tensor, type, tensor_shape, use_device_allocator_for_initializers, alloc)); diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 655dba628b183..b3215c1f2aedb 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -333,9 +333,11 @@ common::Status LoadPrepackedWeightsFromExternalData(const Env& env, const std::f const ONNX_NAMESPACE::TensorProto& tensor_proto, PrepackedWeightsForGraph& prepacked_info); +#if defined(ENABLE_D3D12_FILE_LOADING) common::Status PrepareExtDataForTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, const ONNX_NAMESPACE::TensorProto& tensor_proto, const IExternalDataLoader& ext_data_loader); +#endif // Convert the AttributeProto from a Constant node into a TensorProto that can be used as an initializer // If AttributeProto contains a TensorProto, this tensor proto is converted as is including the case when the diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 48b58c2953d04..d36ba579e7b36 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -499,7 +499,7 @@ struct D3D12AcceleratedExternalDataLoader::Impl { ResolveWeightLoadAccelerationMode( mode, support_status, enabled); if (resolved_status.IsOK() && !enabled && - !IsWeightLoadAccelerationRequired(mode)) { + mode == WeightLoadAccelerationMode::Preferred) { LOGS_DEFAULT(WARNING) << "D3D12 accelerated external weights are unavailable; " "using the ordinary WebGPU initializer loading path. " diff --git a/onnxruntime/test/framework/external_data_loader_test.cc b/onnxruntime/test/framework/external_data_loader_test.cc index 88810be2a6a7d..51aeb4361e4a2 100644 --- a/onnxruntime/test/framework/external_data_loader_test.cc +++ b/onnxruntime/test/framework/external_data_loader_test.cc @@ -308,6 +308,7 @@ TEST_F(ExternalDataLoaderLifetimeTest, RecreatesLoaderAfterFactoryFailure) { #endif +#if defined(ENABLE_D3D12_FILE_LOADING) class BatchLifecycleExternalDataLoader final : public IExternalDataLoader { public: enum class FailurePoint { @@ -401,6 +402,7 @@ TEST(ExternalDataLoaderManagerTest, FinalizeFailureAbortsEveryLoader) { EXPECT_EQ(failing_ptr->abort_count, 1); EXPECT_EQ(unfinalized_ptr->abort_count, 1); } +#endif } // namespace } // namespace test From 2ccf5bfea7e65dcdbcb3fa01f0bb9812485123bb Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:26:09 -0700 Subject: [PATCH 14/46] Complete non-D3D12 loader guards --- onnxruntime/core/framework/session_state_utils.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index f473182a77d1a..9859ba6897a1d 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -343,11 +343,15 @@ common::Status SaveInitializedTensors( // - Values that are external and mapped from disk. We let the OS manage the memory. // - we do not trace values that are in memory because they may be sitting on top of the user allocated // memory. +#if defined(ENABLE_D3D12_FILE_LOADING) const bool loader_creates_tensor = utils::HasExternalData(*tensor_proto) && !utils::HasExternalDataInMemory(*tensor_proto) && external_data_loader_mgr.GetTensorCreator( exec_plan.GetLocation(ort_value_index), tensor_proto->data_type()) != nullptr; +#else + constexpr bool loader_creates_tensor = false; +#endif const bool trace_allocation = !loader_creates_tensor && ((exec_plan.GetLocation(ort_value_index) != default_cpu_device) || !utils::HasExternalData(*tensor_proto)); @@ -369,12 +373,14 @@ common::Status SaveInitializedTensors( // do not trace string tensor continue; } +#if defined(ENABLE_D3D12_FILE_LOADING) if (utils::HasExternalData(*entry.second) && !utils::HasExternalDataInMemory(*entry.second) && external_data_loader_mgr.GetTensorCreator( exec_plan.GetLocation(entry.first), entry.second->data_type()) != nullptr) { continue; } +#endif ORT_RETURN_IF_ERROR(planner.Trace(entry.first, entry.second)); } From f075c3d4064e181db31b6b11e35e6fe234054c69 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:40:37 -0700 Subject: [PATCH 15/46] Normalize D3D12 loading diagnostics and paths --- .../framework/external_data_loader_manager.h | 2 ++ .../webgpu/d3d12_external_data_loader.cc | 17 ++++++++++++++--- .../webgpu/webgpu_provider_factory.cc | 6 +++--- .../providers/webgpu/webgpu_context_test.cc | 19 +++++++++++++++++++ 4 files changed, 38 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index 056d74370e9ca..54c8abf5c5db0 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.h +++ b/onnxruntime/core/framework/external_data_loader_manager.h @@ -3,6 +3,8 @@ #pragma once +#include + #include "core/common/status.h" #include "core/common/common.h" #include "core/framework/external_data_loader.h" diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index d36ba579e7b36..5018a38692ca3 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -49,6 +49,13 @@ double Milliseconds(Clock::duration duration) { return std::chrono::duration(duration).count(); } +struct WindowsPathLess { + bool operator()(const std::filesystem::path& left, + const std::filesystem::path& right) const { + return _wcsicmp(left.c_str(), right.c_str()) < 0; + } +}; + struct TensorKey { std::filesystem::path path; std::string name; @@ -56,8 +63,12 @@ struct TensorKey { size_t length; bool operator<(const TensorKey& other) const { - if (path != other.path) { - return path.native() < other.path.native(); + const WindowsPathLess path_less; + if (path_less(path, other.path)) { + return true; + } + if (path_less(other.path, path)) { + return false; } if (name != other.name) { return name < other.name; @@ -87,7 +98,7 @@ struct D3D12AcceleratedBatch { std::vector tensors; std::map tensors_by_key; - std::map files; + std::map files; size_t total_bytes = 0; size_t load_range_count = 0; bool finalized = false; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index dbac6e8b5b26e..97742385bae5d 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -51,11 +51,11 @@ constexpr const char* WeightLoadAccelerationModeName( WeightLoadAccelerationMode mode) { switch (mode) { case WeightLoadAccelerationMode::Off: - return "off"; + return kWeightLoadAcceleration_Off; case WeightLoadAccelerationMode::Preferred: - return "preferred"; + return kWeightLoadAcceleration_Preferred; case WeightLoadAccelerationMode::Required: - return "required"; + return kWeightLoadAcceleration_Required; } return "unknown"; } diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 50243e6c7885d..4e102d74403d1 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1513,6 +1513,25 @@ TEST(WebGpuContextTest, D3D12AcceleratedLoadsExternalTensorsAcrossFilesAndRanges ASSERT_STATUS_OK(loader->BeginLoad()); ASSERT_STATUS_OK(loader->PrepareTensor( Env::Default(), data_path, "weights", kDataOffset, sizeof(expected))); + auto differently_cased_native = data_path.native(); + ASSERT_FALSE(differently_cased_native.empty()); + if (differently_cased_native.front() >= L'A' && + differently_cased_native.front() <= L'Z') { + differently_cased_native.front() += L'a' - L'A'; + } else { + ASSERT_GE(differently_cased_native.front(), L'a'); + ASSERT_LE(differently_cased_native.front(), L'z'); + differently_cased_native.front() -= L'a' - L'A'; + } + const std::filesystem::path differently_cased_data_path{ + differently_cased_native}; + const auto duplicate_status = loader->PrepareTensor( + Env::Default(), differently_cased_data_path, "weights", + kDataOffset, sizeof(expected)); + EXPECT_FALSE(duplicate_status.IsOK()); + EXPECT_NE(duplicate_status.ErrorMessage().find( + "Duplicate D3D12 accelerated initializer"), + std::string::npos); ASSERT_STATUS_OK(loader->PrepareTensor( Env::Default(), data_path, "separated", kSeparatedDataOffset, sizeof(separated_expected))); From 46b51b5cf52b58de76d44cf78127c7c6388fb871 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:03:06 -0700 Subject: [PATCH 16/46] Harden D3D12 file grouping and cancellation --- .../d3d12_file_buffer_loader.cc | 5 ++- .../webgpu/d3d12_external_data_loader.cc | 35 ++++++++++++------- 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index f91dfb0c15199..9df07ca526d1b 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -552,8 +552,11 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( for (size_t range_index = 0; range_index < ranges.size(); ++range_index) { const auto& range = ranges[range_index]; + auto file_key = range.path; + CharLowerBuffW(file_key.data(), + static_cast(file_key.size())); const auto [file_index_it, inserted] = - file_indices.try_emplace(range.path, files.size()); + file_indices.try_emplace(std::move(file_key), files.size()); if (inserted) { PreparedFile file; file.path = range.path; diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 5018a38692ca3..9d106d3468052 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -114,6 +114,19 @@ struct CancellationState { const std::atomic* abort_requested = nullptr; }; +bool InvokeCancellationNoThrow( + const std::function& is_canceled) noexcept { + if (!is_canceled) { + return false; + } + ORT_TRY { + return is_canceled(); + } + ORT_CATCH(...) { + return true; + } +} + bool IsCancellationRequested(void* opaque) noexcept { const auto& state = *static_cast(opaque); if (state.abort_requested != nullptr && @@ -123,12 +136,7 @@ bool IsCancellationRequested(void* opaque) noexcept { if (state.external == nullptr || !*state.external) { return false; } - ORT_TRY { - return (*state.external)(); - } - ORT_CATCH(...) { - return true; - } + return InvokeCancellationNoThrow(*state.external); } common::Status LoadBatchToD3D12( @@ -138,7 +146,7 @@ common::Status LoadBatchToD3D12( const std::atomic& abort_requested, D3D12AcceleratedLoadMetrics& metrics) { if (abort_requested.load(std::memory_order_relaxed) || - (is_canceled && is_canceled())) { + InvokeCancellationNoThrow(is_canceled)) { return ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, "D3D12 accelerated initializer loading was canceled."); @@ -259,8 +267,9 @@ common::Status PrepareTensorForBatch( ORT_RETURN_IF( batch.tensors_by_key.find(key) != batch.tensors_by_key.end(), - "Duplicate D3D12 accelerated initializer: ", - tensor_name); + "Duplicate D3D12 accelerated initializer \"", + tensor_name, "\" from \"", key.path.string(), + "\" range [", offset, ", ", offset + length, ")."); ORT_RETURN_IF( batch.total_bytes > std::numeric_limits::max() - length, @@ -633,7 +642,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( -> common::Status { if (status.Code() == common::MODEL_LOAD_CANCELED || - (is_canceled && is_canceled())) { + InvokeCancellationNoThrow(is_canceled)) { return ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, "D3D12 accelerated initializer loading was canceled."); @@ -652,7 +661,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( }; if (batch.load_range_count == 0) { - if (is_canceled && is_canceled()) { + if (InvokeCancellationNoThrow(is_canceled)) { return fail_or_fallback(ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, "D3D12 accelerated initializer loading was canceled.")); @@ -673,7 +682,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( return fail_or_fallback(load_status); } - if (is_canceled && is_canceled()) { + if (InvokeCancellationNoThrow(is_canceled)) { return fail_or_fallback(ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, "D3D12 accelerated initializer loading was canceled.")); @@ -687,7 +696,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( "D3D12 accelerated external weights require Dawn " "SharedBufferMemoryD3D12Resource and SharedFenceDXGISharedHandle features."); for (auto& tensor : batch.tensors) { - if (is_canceled && is_canceled()) { + if (InvokeCancellationNoThrow(is_canceled)) { return ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, "D3D12 accelerated initializer loading was canceled."); From 692346bd9b8fc1e40920409f392ff3532bd90cde Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:27:09 -0700 Subject: [PATCH 17/46] Bypass custom loaders for in-memory data --- onnxruntime/core/framework/session_state_utils.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 9859ba6897a1d..268dd1dd4286b 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -99,7 +99,8 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st const auto& memory_info = (alloc != nullptr) ? alloc->Info() : memory_buffer->GetAllocInfo(); const auto device = memory_info.device; - if (utils::HasExternalData(tensor_proto)) { + if (utils::HasExternalData(tensor_proto) && + !utils::HasExternalDataInMemory(tensor_proto)) { auto external_data_loader = external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); if (external_data_loader) { From 5c71fc0144a73c37e322585dfd3ecfd1dad88925 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:49:43 -0700 Subject: [PATCH 18/46] Use ordinal comparison for D3D12 paths --- js/node/test/api/simple-api-tests.ts | 2 +- .../d3d12_file_buffer_loader.cc | 23 +++++++++++++------ .../webgpu/d3d12_external_data_loader.cc | 10 +++++++- 3 files changed, 26 insertions(+), 9 deletions(-) diff --git a/js/node/test/api/simple-api-tests.ts b/js/node/test/api/simple-api-tests.ts index 82a776e50b3e2..f2b33c91fab90 100644 --- a/js/node/test/api/simple-api-tests.ts +++ b/js/node/test/api/simple-api-tests.ts @@ -113,7 +113,7 @@ describe('API Tests - simple API tests', () => { assert.match(message, /weightLoadAcceleration|disk-to-GPU|DXGI|D3D12/); } } else { - assert.equal(message, unavailableProviderError); + assert.doesNotMatch(message, /weightLoadAcceleration|disk-to-GPU|DXGI|D3D12/); } } } diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 9df07ca526d1b..d1341102b93a1 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -9,9 +9,9 @@ #include #include #include +#include #include #include -#include #include #include @@ -34,6 +34,19 @@ constexpr uint64_t kBufferAlignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT constexpr DWORD kInitialCancellationPollMilliseconds = 1; constexpr DWORD kMaximumCancellationPollMilliseconds = 10; +struct WindowsPathLess { + bool operator()(const std::wstring& left, + const std::wstring& right) const noexcept { + const int result = CompareStringOrdinal( + left.data(), static_cast(left.size()), + right.data(), static_cast(right.size()), TRUE); + if (result == 0) { + return left < right; + } + return result == CSTR_LESS_THAN; + } +}; + common::Status HResultError(const char* operation, HRESULT hr) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, operation, " failed with HRESULT 0x", @@ -547,16 +560,12 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( std::vector& files) { // Group ranges by source file so each file is opened once. Later, aligned // overlapping ranges are merged to avoid redundant unbuffered reads. - std::unordered_map file_indices; - file_indices.reserve(ranges.size()); + std::map file_indices; for (size_t range_index = 0; range_index < ranges.size(); ++range_index) { const auto& range = ranges[range_index]; - auto file_key = range.path; - CharLowerBuffW(file_key.data(), - static_cast(file_key.size())); const auto [file_index_it, inserted] = - file_indices.try_emplace(std::move(file_key), files.size()); + file_indices.try_emplace(range.path, files.size()); if (inserted) { PreparedFile file; file.path = range.path; diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 9d106d3468052..9a1257e15757c 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -52,7 +52,15 @@ double Milliseconds(Clock::duration duration) { struct WindowsPathLess { bool operator()(const std::filesystem::path& left, const std::filesystem::path& right) const { - return _wcsicmp(left.c_str(), right.c_str()) < 0; + const auto& left_native = left.native(); + const auto& right_native = right.native(); + const int result = CompareStringOrdinal( + left_native.data(), static_cast(left_native.size()), + right_native.data(), static_cast(right_native.size()), TRUE); + if (result == 0) { + return left_native < right_native; + } + return result == CSTR_LESS_THAN; } }; From aed801e95e6ecb43a2a74c3e5e71066f360e3973 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:12:17 -0700 Subject: [PATCH 19/46] Gate D3D12 import features by loading mode --- .../core/providers/webgpu/webgpu_context.cc | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index db7dc39b4b12d..1a9a8b2f2dad5 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -1109,15 +1109,17 @@ std::vector WebGpuContext::GetAvailableRequiredFeatures(const wgpu::FeatureName::SharedBufferMemoryD3D12Resource, wgpu::FeatureName::SharedFenceDXGISharedHandle, }; - for (auto feature : d3d12_import_features) { - if (adapter.HasFeature(feature)) { - required_features.push_back(feature); - } else { - ORT_ENFORCE(!IsWeightLoadAccelerationRequired( - weight_load_acceleration_mode_), - "The selected Dawn D3D12 adapter does not support a feature required by " - "weightLoadAcceleration: ", - static_cast(feature)); + if (IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_)) { + for (auto feature : d3d12_import_features) { + if (adapter.HasFeature(feature)) { + required_features.push_back(feature); + } else { + ORT_ENFORCE(!IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_), + "The selected Dawn D3D12 adapter does not support a feature required by " + "weightLoadAcceleration: ", + static_cast(feature)); + } } } #endif From 29754587c7fc798dcc79ff52edb9721fea4a4b34 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:01:36 -0700 Subject: [PATCH 20/46] Preserve D3D12 features for shared WebGPU contexts A default WebGPU context can outlive the provider that initialized it in off mode. Request supported D3D12 import features up front so a later accelerated provider can reuse that device, and assert the late-enable contract in the regression test. --- .../core/providers/webgpu/webgpu_context.cc | 24 ++++++++++--------- .../providers/webgpu/webgpu_context_test.cc | 4 ++++ 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 1a9a8b2f2dad5..d1aec80895679 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -1109,17 +1109,19 @@ std::vector WebGpuContext::GetAvailableRequiredFeatures(const wgpu::FeatureName::SharedBufferMemoryD3D12Resource, wgpu::FeatureName::SharedFenceDXGISharedHandle, }; - if (IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_)) { - for (auto feature : d3d12_import_features) { - if (adapter.HasFeature(feature)) { - required_features.push_back(feature); - } else { - ORT_ENFORCE(!IsWeightLoadAccelerationRequired( - weight_load_acceleration_mode_), - "The selected Dawn D3D12 adapter does not support a feature required by " - "weightLoadAcceleration: ", - static_cast(feature)); - } + // The default context can outlive the provider that first initialized it + // (for example, through the shared Env allocator). Request supported import + // features up front so a later provider can enable accelerated loading on + // that same Dawn device. + for (auto feature : d3d12_import_features) { + if (adapter.HasFeature(feature)) { + required_features.push_back(feature); + } else { + ORT_ENFORCE(!IsWeightLoadAccelerationRequired( + weight_load_acceleration_mode_), + "The selected Dawn D3D12 adapter does not support a feature required by " + "weightLoadAcceleration: ", + static_cast(feature)); } } #endif diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 4e102d74403d1..8f698b8e319de 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1292,6 +1292,10 @@ TEST(WebGpuContextTest, D3D12AcceleratedCanBeEnabledAfterInitialOffSession) { if (!support_status.IsOK()) { GTEST_SKIP() << support_status.ErrorMessage(); } + EXPECT_TRUE(context.DeviceHasFeature( + wgpu::FeatureName::SharedBufferMemoryD3D12Resource)); + EXPECT_TRUE(context.DeviceHasFeature( + wgpu::FeatureName::SharedFenceDXGISharedHandle)); auto loader = required_ep->GetExternalDataLoader(); ASSERT_NE(loader, nullptr); From 027a14d0148db592629283fdff9b1cf8aff48e93 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:25:35 -0700 Subject: [PATCH 21/46] Clarify D3D12 loading diagnostics and backend state Include source paths in D3D12 file-open failures and distinguish Dawn's selected backend from the caller's requested backend. --- .../d3d12_file_buffer_loader.cc | 17 +++++++++++++---- .../webgpu/d3d12_external_data_loader.cc | 2 +- .../core/providers/webgpu/webgpu_context.cc | 10 ++++------ .../core/providers/webgpu/webgpu_context.h | 6 +++--- 4 files changed, 21 insertions(+), 14 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index d1341102b93a1..4741579f0fd54 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -60,6 +60,15 @@ common::Status Win32Error(const char* operation, DWORD error) { std::system_category().message(error)); } +common::Status Win32FileError(const char* operation, + const std::wstring& path, + DWORD error) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, operation, " failed for \"", ToUTF8String(path), + "\", error code ", error, " - ", + std::system_category().message(error)); +} + common::Status CancelledStatus() { return ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, @@ -586,8 +595,8 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN, nullptr)); if (!file.file) { - return Win32Error( - "CreateFileW(unbuffered)", GetLastError()); + return Win32FileError( + "CreateFileW(unbuffered)", file.path, GetLastError()); } LARGE_INTEGER file_size{}; @@ -630,8 +639,8 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( FILE_FLAG_SEQUENTIAL_SCAN, nullptr)); if (!file.buffered_tail_file) { - return Win32Error( - "CreateFileW(buffered tail)", GetLastError()); + return Win32FileError( + "CreateFileW(buffered tail)", file.path, GetLastError()); } } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 9a1257e15757c..d5d76aafe8225 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -325,7 +325,7 @@ void EndAccessNoThrow( common::Status CheckD3D12AcceleratedExternalWeightsSupport( const WebGpuContext& context) { ORT_RETURN_IF_NOT( - context.RequestedBackendType() == + context.SelectedBackendType() == wgpu::BackendType::D3D12, "D3D12 accelerated external weights require the Dawn D3D12 backend."); ORT_RETURN_IF_NOT( diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index d1aec80895679..79ba28e1ff0a5 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -85,9 +85,6 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { "adapterIndex requires a native Dawn build with adapter enumeration support."); #endif weight_load_acceleration_mode_ = config.weight_load_acceleration_mode; -#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - requested_backend_type_ = static_cast(config.backend_type); -#endif // Three easily-conflated concepts, at three layers (a pipeline, not the same flag): // * allow_virtual_devices (env) -- selectability: surface a virtual GPU OrtEpDevice so WebGPU is @@ -123,7 +120,8 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { RequestAdapterOptionsLuid luid_options{}; const bool preselect_weight_loading_adapter = IsWeightLoadAccelerationEnabled(weight_load_acceleration_mode_) && - requested_backend_type_ == wgpu::BackendType::D3D12 && + static_cast(config.backend_type) == + wgpu::BackendType::D3D12 && !config.adapter_index; if (preselect_weight_loading_adapter) { LUID selected_luid{}; @@ -420,7 +418,7 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) // Use the backend selected by Dawn rather than the requested backend. The // request can be Undefined when Dawn performs automatic adapter selection. - requested_backend_type_ = adapter_info_.backendType; + selected_backend_type_ = adapter_info_.backendType; #endif // create buffer manager @@ -458,7 +456,7 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) && !defined(USE_EXTERNAL_DAWN) if (IsWeightLoadAccelerationEnabled(config.weight_load_acceleration_mode) && - requested_backend_type_ == wgpu::BackendType::D3D12 && + selected_backend_type_ == wgpu::BackendType::D3D12 && device_ != nullptr && weight_loading_d3d12_device_ == nullptr) { weight_loading_d3d12_device_ = diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index da323f96752fa..a579dfcf36d37 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -341,8 +341,8 @@ class WebGpuContext final { bool D3D12SharedResourceFeaturesAvailable() const { return d3d12_shared_resource_features_available_; } - wgpu::BackendType RequestedBackendType() const { - return requested_backend_type_; + wgpu::BackendType SelectedBackendType() const { + return selected_backend_type_; } #endif @@ -520,7 +520,7 @@ class WebGpuContext final { #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) Microsoft::WRL::ComPtr weight_loading_d3d12_device_; bool d3d12_shared_resource_features_available_ = false; - wgpu::BackendType requested_backend_type_ = wgpu::BackendType::Undefined; + wgpu::BackendType selected_backend_type_ = wgpu::BackendType::Undefined; #endif uint64_t max_storage_buffer_binding_size_; uint32_t max_storage_buffers_per_shader_stage_; From c9573de69adf088464ceb394ef0697f1e3f6ccb8 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:58:50 -0700 Subject: [PATCH 22/46] Clean stale D3D12 loading batches safely Route BeginLoad through the idempotent abort path so imported but unclaimed buffers end Dawn shared-memory access before a new batch replaces them. --- onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index d5d76aafe8225..af2f11e3521c1 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -604,7 +604,7 @@ bool D3D12AcceleratedExternalDataLoader:: common::Status D3D12AcceleratedExternalDataLoader::BeginLoad() const { - impl_->batch.reset(); + AbortLoad(); impl_->ResolveSupport(); ORT_RETURN_IF_ERROR(impl_->resolved_status); if (!impl_->enabled) { From d370fa3655e14e46d843a11b298861003dd93717 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 17:21:22 -0700 Subject: [PATCH 23/46] Export D3D12 loader for static WebGPU installs Install the D3D12 loader archive in the ONNX Runtime CMake export when WebGPU is built statically so the provider's link-only dependency is available to package consumers. --- cmake/onnxruntime_d3d12_file_loader.cmake | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/cmake/onnxruntime_d3d12_file_loader.cmake b/cmake/onnxruntime_d3d12_file_loader.cmake index 8a55047679779..c80d73aeacb90 100644 --- a/cmake/onnxruntime_d3d12_file_loader.cmake +++ b/cmake/onnxruntime_d3d12_file_loader.cmake @@ -26,3 +26,10 @@ target_link_libraries( set_target_properties( onnxruntime_d3d12_file_loader PROPERTIES FOLDER "ONNXRuntime") + +if(NOT onnxruntime_BUILD_SHARED_LIB) + install( + TARGETS onnxruntime_d3d12_file_loader + EXPORT ${PROJECT_NAME}Targets + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endif() From bcd3c5eed0c0618b1c9053ecd5c76369da278b7c Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 17:38:33 -0700 Subject: [PATCH 24/46] Detect D3D12 device removal during fence waits Treat the UINT64_MAX completed-value sentinel as device removal before and after waiting so incomplete copy work cannot be reported as successful. --- .../d3d12_file_buffer_loader.cc | 30 +++++++++++++++---- 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 4741579f0fd54..e96fb1254eb19 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -1047,8 +1047,18 @@ common::Status D3D12FileBufferLoader::Impl::TransitionToCommon( common::Status D3D12FileBufferLoader::Impl::WaitForFence( uint64_t value, const CancellationToken& cancellation) { - if (value == 0 || - copy_fence_->GetCompletedValue() >= value) { + if (value == 0) { + return common::Status::OK(); + } + uint64_t completed_value = + copy_fence_->GetCompletedValue(); + if (completed_value == + std::numeric_limits::max()) { + return HResultError( + "D3D12 device removed while waiting for copy fence", + device_->GetDeviceRemovedReason()); + } + if (completed_value >= value) { return common::Status::OK(); } const HRESULT hr = @@ -1065,14 +1075,22 @@ common::Status D3D12FileBufferLoader::Impl::WaitForFence( const DWORD wait_result = WaitForSingleObject( fence_event_.get(), cancellation_poll_milliseconds); - if (wait_result == WAIT_OBJECT_0) { - return common::Status::OK(); - } - if (wait_result != WAIT_TIMEOUT) { + if (wait_result != WAIT_OBJECT_0 && + wait_result != WAIT_TIMEOUT) { return Win32Error( "WaitForSingleObject(copy fence)", GetLastError()); } + completed_value = copy_fence_->GetCompletedValue(); + if (completed_value == + std::numeric_limits::max()) { + return HResultError( + "D3D12 device removed while waiting for copy fence", + device_->GetDeviceRemovedReason()); + } + if (completed_value >= value) { + return common::Status::OK(); + } if (cancellation.IsCancellationRequested()) { return CancelledStatus(); } From 52dd63cc1affb673d67d52b49bf5dfd88fde5734 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Fri, 2 Oct 2026 17:53:05 -0700 Subject: [PATCH 25/46] Preserve D3D12 file identity and submission safety Keep path casing distinct for Windows case-sensitive directories. Publish fence values only after successful signals, and retain catastrophic untracked submissions for process lifetime rather than hanging or releasing in-flight resources. --- .../d3d12_file_buffer_loader.cc | 93 +++++++++++++------ .../webgpu/d3d12_external_data_loader.cc | 22 +---- .../providers/webgpu/webgpu_context_test.cc | 14 +-- 3 files changed, 69 insertions(+), 60 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index e96fb1254eb19..46dfa1d759b2b 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -34,19 +34,6 @@ constexpr uint64_t kBufferAlignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT constexpr DWORD kInitialCancellationPollMilliseconds = 1; constexpr DWORD kMaximumCancellationPollMilliseconds = 10; -struct WindowsPathLess { - bool operator()(const std::wstring& left, - const std::wstring& right) const noexcept { - const int result = CompareStringOrdinal( - left.data(), static_cast(left.size()), - right.data(), static_cast(right.size()), TRUE); - if (result == 0) { - return left < right; - } - return result == CSTR_LESS_THAN; - } -}; - common::Status HResultError(const char* operation, HRESULT hr) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, operation, " failed with HRESULT 0x", @@ -160,6 +147,8 @@ struct D3D12FileBufferLoader::Impl { const CancellationToken& cancellation); private: + friend class D3D12FileBufferLoader; + common::Status AllocateDestinations( const std::vector& sizes, Batch& batch); @@ -204,8 +193,12 @@ struct D3D12FileBufferLoader::Impl { common::Status WaitForFence( uint64_t value, const CancellationToken& cancellation); + common::Status SignalSubmittedWork(uint64_t value); void WaitForFenceUncancelled(uint64_t value) noexcept; void DrainActiveReads() noexcept; + bool MustRetainUntrackedSubmission() const noexcept { + return retain_untracked_submission_; + } ComPtr device_; Config config_; @@ -213,7 +206,9 @@ struct D3D12FileBufferLoader::Impl { ComPtr copy_fence_; wil::unique_handle fence_event_; std::vector slots_; + Batch untracked_batch_; uint64_t next_fence_value_ = 0; + bool retain_untracked_submission_ = false; std::mutex load_mutex_; }; @@ -334,6 +329,9 @@ common::Status D3D12FileBufferLoader::Impl::Load( std::lock_guard lock(load_mutex_); result.Clear(); + ORT_RETURN_IF( + retain_untracked_submission_, + "D3D12 file buffer loader cannot be reused after an untracked command submission."); if (ranges.empty()) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, @@ -569,7 +567,7 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( std::vector& files) { // Group ranges by source file so each file is opened once. Later, aligned // overlapping ranges are merged to avoid redundant unbuffered reads. - std::map file_indices; + std::map file_indices; for (size_t range_index = 0; range_index < ranges.size(); ++range_index) { const auto& range = ranges[range_index]; @@ -972,14 +970,17 @@ common::Status D3D12FileBufferLoader::Impl::SubmitCopies( slot.command_list.Get()}; copy_queue_->ExecuteCommandLists(1, command_lists); - slot.fence_value = ++next_fence_value_; - last_submitted_fence = slot.fence_value; - hr = copy_queue_->Signal( - copy_fence_.Get(), slot.fence_value); - if (FAILED(hr)) { - return HResultError( - "ID3D12CommandQueue::Signal", hr); + const uint64_t fence_value = ++next_fence_value_; + const auto signal_status = + SignalSubmittedWork(fence_value); + if (!signal_status.IsOK()) { + if (retain_untracked_submission_) { + untracked_batch_ = std::move(batch); + } + return signal_status; } + slot.fence_value = fence_value; + last_submitted_fence = fence_value; return common::Status::OK(); } @@ -1033,17 +1034,45 @@ common::Status D3D12FileBufferLoader::Impl::TransitionToCommon( slot.command_list.Get()}; copy_queue_->ExecuteCommandLists(1, command_lists); - last_submitted_fence = ++next_fence_value_; - hr = copy_queue_->Signal( - copy_fence_.Get(), last_submitted_fence); - if (FAILED(hr)) { - return HResultError( - "ID3D12CommandQueue::Signal(transition)", hr); + const uint64_t fence_value = ++next_fence_value_; + const auto signal_status = + SignalSubmittedWork(fence_value); + if (!signal_status.IsOK()) { + if (retain_untracked_submission_) { + untracked_batch_ = std::move(batch); + } + return signal_status; } + last_submitted_fence = fence_value; return WaitForFence( last_submitted_fence, cancellation); } +common::Status D3D12FileBufferLoader::Impl::SignalSubmittedWork( + uint64_t value) { + // ExecuteCommandLists has no status return. Once work has been submitted, + // its resources must stay alive until either a fence is queued behind it or + // the device is removed. If Signal fails while the device is still live, + // the submission cannot be tracked safely; retain the loader and submitted + // resources for process lifetime rather than hanging or releasing them. + const HRESULT signal_result = + copy_queue_->Signal(copy_fence_.Get(), value); + if (SUCCEEDED(signal_result)) { + return common::Status::OK(); + } + const HRESULT removal_reason = + device_->GetDeviceRemovedReason(); + if (FAILED(removal_reason)) { + return HResultError( + "D3D12 device removed after command submission", + removal_reason); + } + retain_untracked_submission_ = true; + return HResultError( + "ID3D12CommandQueue::Signal after command submission", + signal_result); +} + common::Status D3D12FileBufferLoader::Impl::WaitForFence( uint64_t value, const CancellationToken& cancellation) { @@ -1152,7 +1181,15 @@ D3D12FileBufferLoader::D3D12FileBufferLoader( : impl_(std::move(impl)) { } -D3D12FileBufferLoader::~D3D12FileBufferLoader() = default; +D3D12FileBufferLoader::~D3D12FileBufferLoader() { + if (impl_ && impl_->MustRetainUntrackedSubmission()) { + // A live device has accepted work but rejected the fence that would prove + // its completion. Releasing the referenced resources would be unsafe. + // This catastrophic path intentionally retains the loader for process + // lifetime; normal and device-removed paths destroy it normally. + (void)impl_.release(); + } +} common::Status D3D12FileBufferLoader::Create( ID3D12Device* device, diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index af2f11e3521c1..1d727b3c2b52f 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -49,21 +49,6 @@ double Milliseconds(Clock::duration duration) { return std::chrono::duration(duration).count(); } -struct WindowsPathLess { - bool operator()(const std::filesystem::path& left, - const std::filesystem::path& right) const { - const auto& left_native = left.native(); - const auto& right_native = right.native(); - const int result = CompareStringOrdinal( - left_native.data(), static_cast(left_native.size()), - right_native.data(), static_cast(right_native.size()), TRUE); - if (result == 0) { - return left_native < right_native; - } - return result == CSTR_LESS_THAN; - } -}; - struct TensorKey { std::filesystem::path path; std::string name; @@ -71,11 +56,10 @@ struct TensorKey { size_t length; bool operator<(const TensorKey& other) const { - const WindowsPathLess path_less; - if (path_less(path, other.path)) { + if (path < other.path) { return true; } - if (path_less(other.path, path)) { + if (other.path < path) { return false; } if (name != other.name) { @@ -106,7 +90,7 @@ struct D3D12AcceleratedBatch { std::vector tensors; std::map tensors_by_key; - std::map files; + std::map files; size_t total_bytes = 0; size_t load_range_count = 0; bool finalized = false; diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 8f698b8e319de..3929c7c23e43b 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1517,20 +1517,8 @@ TEST(WebGpuContextTest, D3D12AcceleratedLoadsExternalTensorsAcrossFilesAndRanges ASSERT_STATUS_OK(loader->BeginLoad()); ASSERT_STATUS_OK(loader->PrepareTensor( Env::Default(), data_path, "weights", kDataOffset, sizeof(expected))); - auto differently_cased_native = data_path.native(); - ASSERT_FALSE(differently_cased_native.empty()); - if (differently_cased_native.front() >= L'A' && - differently_cased_native.front() <= L'Z') { - differently_cased_native.front() += L'a' - L'A'; - } else { - ASSERT_GE(differently_cased_native.front(), L'a'); - ASSERT_LE(differently_cased_native.front(), L'z'); - differently_cased_native.front() -= L'a' - L'A'; - } - const std::filesystem::path differently_cased_data_path{ - differently_cased_native}; const auto duplicate_status = loader->PrepareTensor( - Env::Default(), differently_cased_data_path, "weights", + Env::Default(), data_path, "weights", kDataOffset, sizeof(expected)); EXPECT_FALSE(duplicate_status.IsOK()); EXPECT_NE(duplicate_status.ErrorMessage().find( From df3adb8583c9388b7df53ec90af4f340b68fd99a Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 10:36:17 -0700 Subject: [PATCH 26/46] Simplify D3D12 file buffer loader Store loader state and private methods directly on D3D12FileBufferLoader instead of using a separate implementation object. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../d3d12_file_buffer_loader.cc | 221 ++++++------------ .../d3d12_file_buffer_loader.h | 75 +++++- 2 files changed, 145 insertions(+), 151 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 46dfa1d759b2b..00c7288ad311c 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -102,122 +102,68 @@ D3D12_RESOURCE_DESC BufferDescription( } // namespace -struct D3D12FileBufferLoader::Impl { - struct FileReadRegion { - uint64_t begin = 0; - uint64_t end = 0; - }; - - struct PreparedFile { - std::wstring path; - wil::unique_hfile file; - wil::unique_hfile buffered_tail_file; - uint64_t size = 0; - uint64_t alignment = 0; - std::vector range_indices; - std::vector regions; - }; - - struct UploadSlot { - ComPtr resource; - ComPtr allocator; - ComPtr command_list; - wil::unique_handle read_event; - void* mapped = nullptr; - OVERLAPPED overlapped{}; - HANDLE read_file = INVALID_HANDLE_VALUE; - uint64_t file_offset = 0; - DWORD requested = 0; - uint64_t fence_value = 0; - bool read_active = false; - }; +struct D3D12FileBufferLoader::FileReadRegion { + uint64_t begin = 0; + uint64_t end = 0; +}; - struct AllocationResult { - common::Status status; - Batch batch; - }; +struct D3D12FileBufferLoader::PreparedFile { + std::wstring path; + wil::unique_hfile file; + wil::unique_hfile buffered_tail_file; + uint64_t size = 0; + uint64_t alignment = 0; + std::vector range_indices; + std::vector regions; +}; - Impl(ID3D12Device* device, const Config& config); - ~Impl(); - - common::Status Initialize(); - common::Status Load( - const std::vector& ranges, - Batch& result, - const CancellationToken& cancellation); - - private: - friend class D3D12FileBufferLoader; - - common::Status AllocateDestinations( - const std::vector& sizes, - Batch& batch); - common::Status PrepareFiles( - const std::vector& ranges, - const CancellationToken& cancellation, - std::vector& files); - common::Status PrepareSlots( - uint64_t alignment, - const CancellationToken& cancellation); - common::Status IssueRead( - HANDLE file, - UploadSlot& slot, - uint64_t offset, - DWORD size); - common::Status CompleteRead( - uint64_t file_size, - UploadSlot& slot, - DWORD& bytes_read); - template - common::Status ReadRegion( - PreparedFile& file, - const FileReadRegion& region, - const std::vector& ranges, - Batch& batch, - bool& allocation_ready, - EnsureAllocationFn& ensure_allocation, - const CancellationToken& cancellation, - uint64_t& last_submitted_fence); - common::Status SubmitCopies( - UploadSlot& slot, - uint64_t chunk_begin, - uint64_t bytes, - const std::vector& range_indices, - const std::vector& ranges, - Batch& batch, - uint64_t& last_submitted_fence); - common::Status TransitionToCommon( - Batch& batch, - const CancellationToken& cancellation, - uint64_t& last_submitted_fence); - common::Status WaitForFence( - uint64_t value, - const CancellationToken& cancellation); - common::Status SignalSubmittedWork(uint64_t value); - void WaitForFenceUncancelled(uint64_t value) noexcept; - void DrainActiveReads() noexcept; - bool MustRetainUntrackedSubmission() const noexcept { - return retain_untracked_submission_; - } +struct D3D12FileBufferLoader::UploadSlot { + ComPtr resource; + ComPtr allocator; + ComPtr command_list; + wil::unique_handle read_event; + void* mapped = nullptr; + OVERLAPPED overlapped{}; + HANDLE read_file = INVALID_HANDLE_VALUE; + uint64_t file_offset = 0; + DWORD requested = 0; + uint64_t fence_value = 0; + bool read_active = false; +}; - ComPtr device_; - Config config_; - ComPtr copy_queue_; - ComPtr copy_fence_; - wil::unique_handle fence_event_; - std::vector slots_; - Batch untracked_batch_; - uint64_t next_fence_value_ = 0; - bool retain_untracked_submission_ = false; - std::mutex load_mutex_; +struct D3D12FileBufferLoader::AllocationResult { + common::Status status; + Batch batch; }; -D3D12FileBufferLoader::Impl::Impl( - ID3D12Device* device, const Config& config) +D3D12FileBufferLoader::D3D12FileBufferLoader( + ID3D12Device* device, const Config& config) noexcept : device_(device), config_(config) { } -D3D12FileBufferLoader::Impl::~Impl() { +D3D12FileBufferLoader::~D3D12FileBufferLoader() { + if (retain_untracked_submission_) { + // Submitted work has no usable completion fence, so its resources must + // remain alive for the lifetime of the process. + device_.Detach(); + copy_queue_.Detach(); + copy_fence_.Detach(); + (void)fence_event_.release(); + for (auto& slot : slots_) { + slot.resource.Detach(); + slot.allocator.Detach(); + slot.command_list.Detach(); + (void)slot.read_event.release(); + } + for (auto& heap : untracked_batch_.heaps) { + heap.Detach(); + } + for (auto& buffer : untracked_batch_.buffers) { + buffer.resource.Detach(); + } + return; + } + for (auto& slot : slots_) { if (slot.resource && slot.mapped != nullptr) { slot.resource->Unmap(0, nullptr); @@ -226,7 +172,7 @@ D3D12FileBufferLoader::Impl::~Impl() { } } -common::Status D3D12FileBufferLoader::Impl::Initialize() { +common::Status D3D12FileBufferLoader::Initialize() { if (device_ == nullptr) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, @@ -322,7 +268,7 @@ common::Status D3D12FileBufferLoader::Impl::Initialize() { return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::Load( +common::Status D3D12FileBufferLoader::LoadInternal( const std::vector& ranges, Batch& result, const CancellationToken& cancellation) { @@ -447,7 +393,7 @@ common::Status D3D12FileBufferLoader::Impl::Load( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::AllocateDestinations( +common::Status D3D12FileBufferLoader::AllocateDestinations( const std::vector& sizes, Batch& batch) { struct Placement { @@ -561,7 +507,7 @@ common::Status D3D12FileBufferLoader::Impl::AllocateDestinations( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::PrepareFiles( +common::Status D3D12FileBufferLoader::PrepareFiles( const std::vector& ranges, const CancellationToken& cancellation, std::vector& files) { @@ -682,7 +628,7 @@ common::Status D3D12FileBufferLoader::Impl::PrepareFiles( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::PrepareSlots( +common::Status D3D12FileBufferLoader::PrepareSlots( uint64_t alignment, const CancellationToken& cancellation) { for (auto& slot : slots_) { @@ -705,7 +651,7 @@ common::Status D3D12FileBufferLoader::Impl::PrepareSlots( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::IssueRead( +common::Status D3D12FileBufferLoader::IssueRead( HANDLE file, UploadSlot& slot, uint64_t offset, @@ -735,7 +681,7 @@ common::Status D3D12FileBufferLoader::Impl::IssueRead( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::CompleteRead( +common::Status D3D12FileBufferLoader::CompleteRead( uint64_t file_size, UploadSlot& slot, DWORD& bytes_read) { @@ -760,7 +706,7 @@ common::Status D3D12FileBufferLoader::Impl::CompleteRead( } template -common::Status D3D12FileBufferLoader::Impl::ReadRegion( +common::Status D3D12FileBufferLoader::ReadRegion( PreparedFile& file, const FileReadRegion& region, const std::vector& ranges, @@ -894,7 +840,7 @@ common::Status D3D12FileBufferLoader::Impl::ReadRegion( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::SubmitCopies( +common::Status D3D12FileBufferLoader::SubmitCopies( UploadSlot& slot, uint64_t chunk_begin, uint64_t bytes, @@ -984,7 +930,7 @@ common::Status D3D12FileBufferLoader::Impl::SubmitCopies( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::Impl::TransitionToCommon( +common::Status D3D12FileBufferLoader::TransitionToCommon( Batch& batch, const CancellationToken& cancellation, uint64_t& last_submitted_fence) { @@ -1048,7 +994,7 @@ common::Status D3D12FileBufferLoader::Impl::TransitionToCommon( last_submitted_fence, cancellation); } -common::Status D3D12FileBufferLoader::Impl::SignalSubmittedWork( +common::Status D3D12FileBufferLoader::SignalSubmittedWork( uint64_t value) { // ExecuteCommandLists has no status return. Once work has been submitted, // its resources must stay alive until either a fence is queued behind it or @@ -1073,7 +1019,7 @@ common::Status D3D12FileBufferLoader::Impl::SignalSubmittedWork( signal_result); } -common::Status D3D12FileBufferLoader::Impl::WaitForFence( +common::Status D3D12FileBufferLoader::WaitForFence( uint64_t value, const CancellationToken& cancellation) { if (value == 0) { @@ -1130,7 +1076,7 @@ common::Status D3D12FileBufferLoader::Impl::WaitForFence( } } -void D3D12FileBufferLoader::Impl::WaitForFenceUncancelled( +void D3D12FileBufferLoader::WaitForFenceUncancelled( uint64_t value) noexcept { if (value == 0 || copy_fence_->GetCompletedValue() >= value) { @@ -1155,7 +1101,7 @@ void D3D12FileBufferLoader::Impl::WaitForFenceUncancelled( } } -void D3D12FileBufferLoader::Impl::DrainActiveReads() noexcept { +void D3D12FileBufferLoader::DrainActiveReads() noexcept { for (auto& slot : slots_) { if (!slot.read_active) { continue; @@ -1176,31 +1122,16 @@ void D3D12FileBufferLoader::Impl::DrainActiveReads() noexcept { } } -D3D12FileBufferLoader::D3D12FileBufferLoader( - std::unique_ptr impl) noexcept - : impl_(std::move(impl)) { -} - -D3D12FileBufferLoader::~D3D12FileBufferLoader() { - if (impl_ && impl_->MustRetainUntrackedSubmission()) { - // A live device has accepted work but rejected the fence that would prove - // its completion. Releasing the referenced resources would be unsafe. - // This catastrophic path intentionally retains the loader for process - // lifetime; normal and device-removed paths destroy it normally. - (void)impl_.release(); - } -} - common::Status D3D12FileBufferLoader::Create( ID3D12Device* device, std::unique_ptr& loader, const Config& config) noexcept { loader.reset(); try { - auto impl = std::make_unique(device, config); - ORT_RETURN_IF_ERROR(impl->Initialize()); - loader.reset( - new D3D12FileBufferLoader(std::move(impl))); + auto instance = std::unique_ptr( + new D3D12FileBufferLoader(device, config)); + ORT_RETURN_IF_ERROR(instance->Initialize()); + loader = std::move(instance); return common::Status::OK(); } catch (const std::exception& ex) { return ORT_MAKE_STATUS( @@ -1218,14 +1149,8 @@ common::Status D3D12FileBufferLoader::Load( const std::vector& ranges, Batch& result, const CancellationToken& cancellation) noexcept { - if (!impl_) { - result.Clear(); - return ORT_MAKE_STATUS( - ONNXRUNTIME, FAIL, - "D3D12FileBufferLoader is not initialized."); - } try { - return impl_->Load(ranges, result, cancellation); + return LoadInternal(ranges, result, cancellation); } catch (const std::exception& ex) { result.Clear(); return ORT_MAKE_STATUS( diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h index 13d4cfb208b5b..25efdb12e975d 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h @@ -5,10 +5,12 @@ #include #include +#include #include #include #include +#include #include #include "core/common/status.h" @@ -76,11 +78,78 @@ class D3D12FileBufferLoader { const CancellationToken& cancellation = {}) noexcept; private: - struct Impl; + struct FileReadRegion; + struct PreparedFile; + struct UploadSlot; + struct AllocationResult; - explicit D3D12FileBufferLoader(std::unique_ptr impl) noexcept; + D3D12FileBufferLoader( + ID3D12Device* device, + const Config& config) noexcept; - std::unique_ptr impl_; + common::Status Initialize(); + common::Status LoadInternal( + const std::vector& ranges, + Batch& result, + const CancellationToken& cancellation); + common::Status AllocateDestinations( + const std::vector& sizes, + Batch& batch); + common::Status PrepareFiles( + const std::vector& ranges, + const CancellationToken& cancellation, + std::vector& files); + common::Status PrepareSlots( + uint64_t alignment, + const CancellationToken& cancellation); + common::Status IssueRead( + HANDLE file, + UploadSlot& slot, + uint64_t offset, + DWORD size); + common::Status CompleteRead( + uint64_t file_size, + UploadSlot& slot, + DWORD& bytes_read); + template + common::Status ReadRegion( + PreparedFile& file, + const FileReadRegion& region, + const std::vector& ranges, + Batch& batch, + bool& allocation_ready, + EnsureAllocationFn& ensure_allocation, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence); + common::Status SubmitCopies( + UploadSlot& slot, + uint64_t chunk_begin, + uint64_t bytes, + const std::vector& range_indices, + const std::vector& ranges, + Batch& batch, + uint64_t& last_submitted_fence); + common::Status TransitionToCommon( + Batch& batch, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence); + common::Status WaitForFence( + uint64_t value, + const CancellationToken& cancellation); + common::Status SignalSubmittedWork(uint64_t value); + void WaitForFenceUncancelled(uint64_t value) noexcept; + void DrainActiveReads() noexcept; + + Microsoft::WRL::ComPtr device_; + Config config_; + Microsoft::WRL::ComPtr copy_queue_; + Microsoft::WRL::ComPtr copy_fence_; + wil::unique_handle fence_event_; + std::vector slots_; + Batch untracked_batch_; + uint64_t next_fence_value_ = 0; + bool retain_untracked_submission_ = false; + std::mutex load_mutex_; }; } // namespace d3d12 From 9ac9de2ba587cad904fddaa1c678f77e2f45b827 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 11:23:49 -0700 Subject: [PATCH 27/46] Handle custom WebGPU initializer allocators Select tensor-creating external data loaders using the active initializer allocator. Preferred D3D12 loading falls back to ordinary loading for incompatible allocators, while required mode reports an incompatibility error. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/framework/external_data_loader.h | 10 ++ .../framework/external_data_loader_manager.cc | 55 ++++++++-- .../framework/external_data_loader_manager.h | 12 ++- .../core/framework/session_state_utils.cc | 40 +++++-- .../core/providers/webgpu/allocator.cc | 7 +- .../webgpu/d3d12_external_data_loader.cc | 48 +++++++-- .../webgpu/d3d12_external_data_loader.h | 9 ++ .../core/providers/webgpu/webgpu_context.cc | 8 +- .../core/providers/webgpu/webgpu_context.h | 1 + .../webgpu/concurrent_context_test.cc | 100 ++++++++++++++++++ .../providers/webgpu/webgpu_context_test.cc | 45 +++++++- 11 files changed, 296 insertions(+), 39 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index b935a028a3ce6..86bbbbfc3a17f 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -42,6 +42,16 @@ class IExternalDataLoader { // writing into a tensor allocated by the framework. virtual bool CreatesTensorForDevice([[maybe_unused]] const OrtDevice& target_device) const { return false; } + // Determines whether a tensor-creating loader can use the allocator selected for + // the SessionState. An OK status with can_create_tensor=false lets ordinary + // initializer loading handle the tensor. + virtual common::Status CanCreateTensorWithAllocator( + [[maybe_unused]] const std::shared_ptr& allocator, + bool& can_create_tensor) const { + can_create_tensor = true; + return common::Status::OK(); + } + // Begins a batch covering the external initializers for one SessionState // initialization. ORT describes all loader-created tensors through PrepareTensor() // before calling FinalizeLoad(), allowing an implementation to coalesce file diff --git a/onnxruntime/core/framework/external_data_loader_manager.cc b/onnxruntime/core/framework/external_data_loader_manager.cc index 86746b4a2126c..26ecf71e4254c 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.cc +++ b/onnxruntime/core/framework/external_data_loader_manager.cc @@ -3,6 +3,8 @@ #include "core/framework/external_data_loader_manager.h" +#include "core/framework/allocator.h" + namespace onnxruntime { using namespace common; @@ -37,16 +39,57 @@ const IExternalDataLoader* ExternalDataLoaderManager::GetExternalDataLoader( } #if defined(ENABLE_D3D12_FILE_LOADING) -const IExternalDataLoader* ExternalDataLoaderManager::GetTensorCreator( - const OrtDevice& target_device, int32_t tensor_data_type) const { +Status ExternalDataLoaderManager::GetTensorCreator( + const OrtDevice& target_device, + int32_t tensor_data_type, + const std::shared_ptr& allocator, + const IExternalDataLoader*& tensor_creator) const { + tensor_creator = nullptr; for (const auto& external_data_loader : external_data_loaders_) { - if (external_data_loader->SupportsDataType(tensor_data_type) && - external_data_loader->CreatesTensorForDevice(target_device)) { - return external_data_loader.get(); + if (!external_data_loader->SupportsDataType(tensor_data_type) || + !external_data_loader->CreatesTensorForDevice(target_device)) { + continue; + } + + bool can_create_tensor = false; + ORT_RETURN_IF_ERROR( + external_data_loader->CanCreateTensorWithAllocator(allocator, can_create_tensor)); + if (can_create_tensor) { + tensor_creator = external_data_loader.get(); + break; } } - return nullptr; + return Status::OK(); +} + +Status ExternalDataLoaderManager::GetExternalDataLoader( + const OrtMemoryInfo& target_memory_info, + int32_t tensor_data_type, + const std::shared_ptr& allocator, + const IExternalDataLoader*& selected_loader) const { + selected_loader = nullptr; + for (const auto& external_data_loader : external_data_loaders_) { + if (!external_data_loader->SupportsDataType(tensor_data_type)) { + continue; + } + + if (external_data_loader->CreatesTensorForDevice(target_memory_info.device)) { + bool can_create_tensor = false; + ORT_RETURN_IF_ERROR( + external_data_loader->CanCreateTensorWithAllocator(allocator, can_create_tensor)); + if (!can_create_tensor) { + continue; + } + } else if (!external_data_loader->CanLoad(target_memory_info)) { + continue; + } + + selected_loader = external_data_loader.get(); + break; + } + + return Status::OK(); } Status ExternalDataLoaderManager::BeginLoad() const { diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index 54c8abf5c5db0..5a5df43df8fdd 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.h +++ b/onnxruntime/core/framework/external_data_loader_manager.h @@ -27,8 +27,16 @@ class ExternalDataLoaderManager { const IExternalDataLoader* GetExternalDataLoader( const OrtMemoryInfo& target_memory_info, int32_t tensor_data_type) const; #if defined(ENABLE_D3D12_FILE_LOADING) - const IExternalDataLoader* GetTensorCreator( - const OrtDevice& target_device, int32_t tensor_data_type) const; + common::Status GetTensorCreator( + const OrtDevice& target_device, + int32_t tensor_data_type, + const std::shared_ptr& allocator, + const IExternalDataLoader*& tensor_creator) const; + common::Status GetExternalDataLoader( + const OrtMemoryInfo& target_memory_info, + int32_t tensor_data_type, + const std::shared_ptr& allocator, + const IExternalDataLoader*& external_data_loader) const; common::Status BeginLoad() const; common::Status FinalizeLoad(const std::function& is_canceled) const; void AbortLoad() const noexcept; diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 268dd1dd4286b..e5dc3d1475d05 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -101,8 +101,14 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st if (utils::HasExternalData(tensor_proto) && !utils::HasExternalDataInMemory(tensor_proto)) { +#if defined(ENABLE_D3D12_FILE_LOADING) + const IExternalDataLoader* external_data_loader = nullptr; + ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetExternalDataLoader( + memory_info, tensor_proto.data_type(), alloc, external_data_loader)); +#else auto external_data_loader = external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); +#endif if (external_data_loader) { #if defined(ENABLE_D3D12_FILE_LOADING) if (external_data_loader->CreatesTensorForDevice(device)) { @@ -345,11 +351,15 @@ common::Status SaveInitializedTensors( // - we do not trace values that are in memory because they may be sitting on top of the user allocated // memory. #if defined(ENABLE_D3D12_FILE_LOADING) - const bool loader_creates_tensor = - utils::HasExternalData(*tensor_proto) && - !utils::HasExternalDataInMemory(*tensor_proto) && - external_data_loader_mgr.GetTensorCreator( - exec_plan.GetLocation(ort_value_index), tensor_proto->data_type()) != nullptr; + const IExternalDataLoader* tensor_creator = nullptr; + if (utils::HasExternalData(*tensor_proto) && + !utils::HasExternalDataInMemory(*tensor_proto)) { + const auto& target_device = exec_plan.GetLocation(ort_value_index); + ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetTensorCreator( + target_device, tensor_proto->data_type(), + planner.GetInitializerAllocator(target_device), tensor_creator)); + } + const bool loader_creates_tensor = tensor_creator != nullptr; #else constexpr bool loader_creates_tensor = false; #endif @@ -376,10 +386,15 @@ common::Status SaveInitializedTensors( } #if defined(ENABLE_D3D12_FILE_LOADING) if (utils::HasExternalData(*entry.second) && - !utils::HasExternalDataInMemory(*entry.second) && - external_data_loader_mgr.GetTensorCreator( - exec_plan.GetLocation(entry.first), entry.second->data_type()) != nullptr) { - continue; + !utils::HasExternalDataInMemory(*entry.second)) { + const auto& target_device = exec_plan.GetLocation(entry.first); + const IExternalDataLoader* tensor_creator = nullptr; + ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetTensorCreator( + target_device, entry.second->data_type(), + planner.GetInitializerAllocator(target_device), tensor_creator)); + if (tensor_creator != nullptr) { + continue; + } } #endif ORT_RETURN_IF_ERROR(planner.Trace(entry.first, entry.second)); @@ -428,8 +443,11 @@ common::Status SaveInitializedTensors( "Preparing session state weights is canceled due to user request."); } - const auto* tensor_creator = external_data_loader_mgr.GetTensorCreator( - exec_plan.GetLocation(entry.first), entry.second->data_type()); + const auto& target_device = exec_plan.GetLocation(entry.first); + const IExternalDataLoader* tensor_creator = nullptr; + ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetTensorCreator( + target_device, entry.second->data_type(), + planner.GetInitializerAllocator(target_device), tensor_creator)); if (tensor_creator != nullptr) { ORT_RETURN_IF_ERROR( utils::PrepareExtDataForTensorFromTensorProto(env, graph_loc, *entry.second, *tensor_creator)); diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index fd210c5b5896e..a697a6ae8e811 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -50,11 +50,8 @@ void* GpuBufferAllocator::Allocate(size_t size, CommandRecordingState& recording } const auto& buffer_manager = buffer_manager_getter_(); - if (!mapped_at_creation_.has_value()) { - mapped_at_creation_ = is_read_only_allocator_ && buffer_manager.SupportsUMA(); - } - wgpu::BufferUsage usage = *mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite - : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; + wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite + : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; auto buffer = buffer_manager.Create(recording, size, usage, initialize_to_zero_, submit_zero_initialize); num_allocs_.fetch_add(1, std::memory_order_relaxed); diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 1d727b3c2b52f..6fee42297ad69 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -407,8 +407,6 @@ class D3D12AcceleratedWebGpuAllocator final wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; auto& recording = recording_getter_(); - std::lock_guard lock{ - recording.mutex}; return context_.InitializerBufferManager().Create( recording, size, usage); } @@ -419,8 +417,6 @@ class D3D12AcceleratedWebGpuAllocator final } auto& recording = recording_getter_(); - std::lock_guard - recording_lock{recording.mutex}; std::unique_ptr imported; { std::lock_guard lock{ @@ -437,17 +433,26 @@ class D3D12AcceleratedWebGpuAllocator final } if (imported) { - if (recording.has_unsubmitted_work) { + if (recording.has_unsubmitted_work.load( + std::memory_order_acquire)) { std::shared_ptr deferred{ imported.release(), [](ImportedAllocation* allocation) { EndAccessNoThrow(*allocation); delete allocation; }}; - recording.pending_release_callbacks.emplace_back( - [deferred = std::move(deferred)]() mutable { - deferred.reset(); - }); + { + std::lock_guard lock{ + recording.pending_release_callbacks_mutex}; + if (recording.has_unsubmitted_work.load( + std::memory_order_acquire)) { + recording.pending_release_callbacks.emplace_back( + [deferred]() mutable { + deferred.reset(); + }); + } + } + deferred.reset(); } else { EndAccessNoThrow(*imported); } @@ -455,7 +460,7 @@ class D3D12AcceleratedWebGpuAllocator final } context_.InitializerBufferManager().Release( - recording, static_cast(p)); + static_cast(p), &recording); } private: @@ -586,6 +591,29 @@ bool D3D12AcceleratedExternalDataLoader:: IsWeightLoadAccelerationRequired(impl_->mode)); } +common::Status ResolveWeightLoadAccelerationAllocator( + WeightLoadAccelerationMode mode, + const IAllocator* accelerated_allocator, + const std::shared_ptr& selected_allocator, + bool& compatible) { + compatible = selected_allocator != nullptr && + selected_allocator.get() == accelerated_allocator; + ORT_RETURN_IF( + !compatible && IsWeightLoadAccelerationRequired(mode), + "weightLoadAcceleration=\"required\" is incompatible with the selected " + "WebGPU initializer allocator."); + return common::Status::OK(); +} + +common::Status D3D12AcceleratedExternalDataLoader:: + CanCreateTensorWithAllocator( + const std::shared_ptr& allocator, + bool& can_create_tensor) const { + return ResolveWeightLoadAccelerationAllocator( + impl_->mode, impl_->state->impl_->allocator, + allocator, can_create_tensor); +} + common::Status D3D12AcceleratedExternalDataLoader::BeginLoad() const { AbortLoad(); diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index 7f10d63193f24..7f4eadf7a10b3 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -28,6 +28,12 @@ common::Status ResolveWeightLoadAccelerationMode( const common::Status& support_status, bool& enabled); +common::Status ResolveWeightLoadAccelerationAllocator( + WeightLoadAccelerationMode mode, + const IAllocator* accelerated_allocator, + const std::shared_ptr& selected_allocator, + bool& compatible); + // Shared by the loader and allocator so imported resources outlive initializer loading. class D3D12AcceleratedInitializerState { public: @@ -64,6 +70,9 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { bool CanLoad(const OrtMemoryInfo& target_memory_info) const override; bool SupportsDataType(int32_t tensor_data_type) const override; bool CreatesTensorForDevice(const OrtDevice& target_device) const override; + common::Status CanCreateTensorWithAllocator( + const std::shared_ptr& allocator, + bool& can_create_tensor) const override; common::Status BeginLoad() const override; common::Status PrepareTensor(const Env& env, const std::filesystem::path& data_file_path, diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 79ba28e1ff0a5..664c375ea023c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -1453,10 +1453,14 @@ Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, device_queue_.Submit(1, &command_buffer); refresh_pending_buffers(); #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) - for (auto& release : recording.pending_release_callbacks) { + std::vector> pending_release_callbacks; + { + std::lock_guard lock{recording.pending_release_callbacks_mutex}; + pending_release_callbacks.swap(recording.pending_release_callbacks); + } + for (auto& release : pending_release_callbacks) { release(); } - recording.pending_release_callbacks.clear(); #endif recording.command_encoder = nullptr; recording.num_pending_dispatches = 0; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index a579dfcf36d37..e02777b1d78b3 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -130,6 +130,7 @@ struct CommandRecordingState { // BufferManager clears it under its cache lock after submission or abandonment. std::atomic has_unsubmitted_work{false}; #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + std::mutex pending_release_callbacks_mutex; std::vector> pending_release_callbacks; #endif std::vector deferred_dispatches; diff --git a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc index a2783095a16e0..0118613df8c29 100644 --- a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc +++ b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include #include #include #include @@ -8,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -19,7 +22,11 @@ #include "core/graph/model.h" #include "core/graph/onnx_protobuf.h" #include "core/platform/env.h" +#include "core/providers/webgpu/allocator.h" +#include "core/providers/webgpu/data_transfer.h" #include "core/providers/webgpu/program_manager.h" +#include "core/providers/webgpu/webgpu_context.h" +#include "core/providers/webgpu/webgpu_external_header.h" #include "core/providers/webgpu/webgpu_provider_options.h" #include "core/session/inference_session.h" @@ -31,6 +38,98 @@ namespace onnxruntime { namespace test { +namespace { + +void BuildAddChainModel(int chain_len, int64_t num_elements, std::string& model_bytes) { + const std::unordered_map domain_to_version{{"", 13}}; + Model model("webgpu_concurrent_ctx", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector(), + DefaultLoggingManager().DefaultLogger()); + Graph& graph = model.MainGraph(); + + ONNX_NAMESPACE::TypeProto float_1d; + float_1d.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + float_1d.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(num_elements); + + const std::vector weight_values(static_cast(num_elements), 0.5f); + + NodeArg* prev = &graph.GetOrCreateNodeArg("X", &float_1d); + for (int i = 0; i < chain_len; ++i) { + const std::string w_name = "W" + std::to_string(i); + ONNX_NAMESPACE::TensorProto w_tensor; + w_tensor.set_name(w_name); + w_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + w_tensor.add_dims(num_elements); + w_tensor.set_raw_data(weight_values.data(), weight_values.size() * sizeof(float)); + graph.AddInitializedTensor(w_tensor); + + NodeArg* w_arg = &graph.GetOrCreateNodeArg(w_name, &float_1d); + const std::string out_name = (i == chain_len - 1) ? "Y" : ("H" + std::to_string(i)); + NodeArg* out_arg = &graph.GetOrCreateNodeArg(out_name, &float_1d); + graph.AddNode("add" + std::to_string(i), "Add", "", {prev, w_arg}, {out_arg}); + prev = out_arg; + } + + graph.SetOutputs({prev}); + ASSERT_STATUS_OK(graph.Resolve()); + ASSERT_TRUE(model.ToProto().SerializeToString(&model_bytes)); +} + +constexpr const char* kUnaryOps[] = { + "Abs", "Neg", "Floor", "Ceil", "Reciprocal", "Sqrt", "Exp", "Erf", "Sigmoid", + "Sin", "Cos", "Tan", "Atan", "Sinh", "Cosh", "Tanh", "HardSigmoid", "HardSwish"}; + +void BuildUnaryFanOutModel(int64_t num_elements, std::string& model_bytes) { + const std::unordered_map domain_to_version{{"", 14}}; + Model model("webgpu_concurrent_ctx_cold", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector(), + DefaultLoggingManager().DefaultLogger()); + Graph& graph = model.MainGraph(); + + ONNX_NAMESPACE::TypeProto float_1d; + float_1d.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + float_1d.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(num_elements); + + NodeArg* x_arg = &graph.GetOrCreateNodeArg("X", &float_1d); + std::vector outputs; + outputs.reserve(std::size(kUnaryOps)); + for (size_t i = 0; i < std::size(kUnaryOps); ++i) { + NodeArg* out_arg = &graph.GetOrCreateNodeArg("Y" + std::to_string(i), &float_1d); + graph.AddNode("op" + std::to_string(i), kUnaryOps[i], "", {x_arg}, {out_arg}); + outputs.push_back(out_arg); + } + + graph.SetOutputs(outputs); + ASSERT_STATUS_OK(graph.Resolve()); + ASSERT_TRUE(model.ToProto().SerializeToString(&model_bytes)); +} + +class ErrorSink { + public: + void Record(const std::string& message) { + bool expected = false; + if (failed_.compare_exchange_strong(expected, true)) { + std::lock_guard lock(mutex_); + first_error_ = message; + } + } + + bool Failed() const { + return failed_.load(); + } + + std::string FirstError() const { + std::lock_guard lock(mutex_); + return first_error_; + } + + private: + std::atomic failed_{false}; + mutable std::mutex mutex_; + std::string first_error_; +}; TEST(WebGpuConcurrentContextTestStandalone, ShaderDumpWritesRemainComplete) { TemporaryDirectory temp_dir{ORT_TSTR("webgpu_shader_dump_test")}; @@ -646,5 +745,6 @@ TEST(WebGpuPoolMemory, DISABLED_MultiSessionSameShape) { std::this_thread::sleep_for(std::chrono::milliseconds(hold_ms)); } +} // namespace } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 3929c7c23e43b..5fe9ebb7deac3 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1422,11 +1422,22 @@ TEST(WebGpuContextTest, D3D12AcceleratedAllocatorUsesProviderRecording) { auto& recording = webgpu_ep.Recording(); recording.has_unsubmitted_work = true; allocator->Free(buffer); - EXPECT_EQ(recording.pending_buffers.size(), 1u); + constexpr auto usage = + wgpu::BufferUsage::Storage | + wgpu::BufferUsage::CopySrc | + wgpu::BufferUsage::CopyDst | + wgpu::BufferUsage::Indirect; + auto& buffer_manager = webgpu_ep.InitializerBufferManager(); + EXPECT_TRUE(buffer_manager.ExtractCachedBuffers(usage).empty()); recording.has_unsubmitted_work = false; - webgpu_ep.InitializerBufferManager().RefreshPendingBuffers(recording); - EXPECT_TRUE(recording.pending_buffers.empty()); + buffer_manager.RefreshPendingBuffers( + recording, webgpu::GraphCaptureState::Default); + auto cached_buffers = + buffer_manager.ExtractCachedBuffers(usage); + EXPECT_EQ(cached_buffers.size(), 1u); + buffer_manager.AbsorbCachedBuffers( + usage, std::move(cached_buffers)); } TEST(WebGpuContextTest, DuplicateProviderDoesNotRegisterD3D12AcceleratedLoader) { @@ -1922,6 +1933,34 @@ TEST(WebGpuContextTest, WeightLoadAccelerationModeResolution) { EXPECT_FALSE(enabled); EXPECT_EQ(required_status.ErrorMessage(), unsupported.ErrorMessage()); } + +TEST(WebGpuContextTest, WeightLoadAccelerationAllocatorResolution) { + const auto accelerated_allocator = std::make_shared(); + const auto environment_allocator = std::make_shared(); + bool compatible = false; + + EXPECT_STATUS_OK(webgpu::ResolveWeightLoadAccelerationAllocator( + webgpu::WeightLoadAccelerationMode::Preferred, + accelerated_allocator.get(), accelerated_allocator, compatible)); + EXPECT_TRUE(compatible); + + compatible = true; + EXPECT_STATUS_OK(webgpu::ResolveWeightLoadAccelerationAllocator( + webgpu::WeightLoadAccelerationMode::Preferred, + accelerated_allocator.get(), environment_allocator, compatible)); + EXPECT_FALSE(compatible); + + compatible = true; + const auto required_status = + webgpu::ResolveWeightLoadAccelerationAllocator( + webgpu::WeightLoadAccelerationMode::Required, + accelerated_allocator.get(), environment_allocator, compatible); + EXPECT_FALSE(required_status.IsOK()); + EXPECT_FALSE(compatible); + EXPECT_NE(required_status.ErrorMessage().find( + "incompatible with the selected WebGPU initializer allocator"), + std::string::npos); +} #endif TEST(WebGpuContextTest, AdapterIndexRejectsConflictingSelectorOnReusedContext) { From 0136742648c8283be009d0451c0f2a8ace48adff Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:17:51 -0700 Subject: [PATCH 28/46] Serialize shared WebGPU command recording Protect command recording state shared by runs, allocators, data transfers, graph capture, and accelerated initializer allocation. This prevents concurrent encoder mutation while retaining the current pending-buffer synchronization. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/providers/webgpu/allocator.cc | 10 ++++++++-- .../providers/webgpu/d3d12_external_data_loader.cc | 4 ++++ onnxruntime/core/providers/webgpu/data_transfer.cc | 3 +++ onnxruntime/core/providers/webgpu/webgpu_context.cc | 8 +++++++- onnxruntime/core/providers/webgpu/webgpu_context.h | 6 +++--- 5 files changed, 25 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index a697a6ae8e811..dc973c801f7aa 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include @@ -49,6 +50,7 @@ void* GpuBufferAllocator::Allocate(size_t size, CommandRecordingState& recording return nullptr; } + std::lock_guard lock{recording.mutex}; const auto& buffer_manager = buffer_manager_getter_(); wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; @@ -60,8 +62,12 @@ void* GpuBufferAllocator::Allocate(size_t size, CommandRecordingState& recording void GpuBufferAllocator::Free(void* p) { if (p != nullptr) { - buffer_manager_getter_().Release(static_cast(p), - recording_getter_ ? &recording_getter_() : nullptr); + auto* recording = recording_getter_ ? &recording_getter_() : nullptr; + std::unique_lock lock; + if (recording != nullptr) { + lock = std::unique_lock{recording->mutex}; + } + buffer_manager_getter_().Release(static_cast(p), recording); num_allocs_.fetch_sub(1, std::memory_order_relaxed); } } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 6fee42297ad69..2ff73c2f137e5 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -407,6 +407,8 @@ class D3D12AcceleratedWebGpuAllocator final wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; auto& recording = recording_getter_(); + std::lock_guard lock{ + recording.mutex}; return context_.InitializerBufferManager().Create( recording, size, usage); } @@ -417,6 +419,8 @@ class D3D12AcceleratedWebGpuAllocator final } auto& recording = recording_getter_(); + std::lock_guard + recording_lock{recording.mutex}; std::unique_ptr imported; { std::lock_guard lock{ diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 4411c922a6fc3..4b8ec7c523043 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -3,6 +3,8 @@ #include "core/providers/webgpu/data_transfer.h" +#include + #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/webgpu_context.h" @@ -15,6 +17,7 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data, bool dst_is_gpu, size_t bytes) const { auto& command_state = recording_; + std::lock_guard lock{command_state.mutex}; if (bytes > 0) { if (dst_is_gpu) { if (src_is_gpu) { diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 664c375ea023c..0643654634da5 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -632,6 +632,7 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra const webgpu::BufferManager& buffer_mgr = ComputeContextBase::BufferManagerAccessor::Get(context); CommandRecordingState& recording = ComputeContextBase::BufferManagerAccessor::GetRecording(context); + std::lock_guard lock{recording.mutex}; // validate inputs and outputs are on WebGPU buffers if (ValidationMode() >= ValidationMode::Basic) { @@ -1390,6 +1391,7 @@ Status WebGpuContext::FlushAndWaitChecked(const webgpu::BufferManager& buffer_mg Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording) { + std::lock_guard lock{recording.mutex}; // Graph runs and prepacking can release allocator/uniform buffers into the shared manager. // Retire this recording in both managers, keeping graph cache policy local to the active one. const auto refresh_pending_buffers = [&]() { @@ -1455,7 +1457,8 @@ Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) std::vector> pending_release_callbacks; { - std::lock_guard lock{recording.pending_release_callbacks_mutex}; + std::lock_guard callbacks_lock{ + recording.pending_release_callbacks_mutex}; pending_release_callbacks.swap(recording.pending_release_callbacks); } for (auto& release : pending_release_callbacks) { @@ -1536,6 +1539,7 @@ void WebGpuContext::DispatchCommand(const webgpu::CapturedCommandInfo& command, void WebGpuContext::CaptureBegin(std::vector* captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { + std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "CaptureBegin with external storage"; // Flush any pending commands before we change the status ORT_THROW_IF_ERROR(Flush(buffer_manager, recording)); @@ -1552,6 +1556,7 @@ void WebGpuContext::CaptureBegin(std::vector* captu void WebGpuContext::Replay(const std::vector& captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { + std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "Replay with external storage"; recording.graph_capture_state = GraphCaptureState::Replaying; // Replay all captured commands from the provided vector @@ -1582,6 +1587,7 @@ void WebGpuContext::Replay(const std::vector& captu } void WebGpuContext::CaptureEnd(CommandRecordingState& recording) { + std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "CaptureEnd"; recording.graph_capture_state = GraphCaptureState::Default; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index e02777b1d78b3..423d5229f11f2 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -116,13 +116,13 @@ struct CapturedCommandInfo { std::optional pending_kernel_info; }; -// A command recording timeline has one caller at a time. ORT serializes Run on a Session; -// plugin streamless operations use independent encoders. BufferManager synchronizes buffer -// releases from concurrent allocator callers separately from command recording. +// Serializes callers that share a command recording timeline, including Run, public +// allocator use, data transfer, and graph capture operations. struct CommandRecordingState { CommandRecordingState() = default; ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CommandRecordingState); + std::recursive_mutex mutex; wgpu::CommandEncoder command_encoder; wgpu::ComputePassEncoder compute_pass_encoder; uint32_t num_pending_dispatches = 0; From 4a7c942d4b0afc5893df0f99631f3428602f77c4 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:34:49 -0700 Subject: [PATCH 29/46] Simplify WebGPU allocator compatibility Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/framework/external_data_loader.h | 10 ---- .../framework/external_data_loader_manager.cc | 55 ++--------------- .../framework/external_data_loader_manager.h | 12 +--- .../core/framework/session_state_utils.cc | 40 ++++--------- .../webgpu/d3d12_external_data_loader.cc | 23 ------- .../webgpu/d3d12_external_data_loader.h | 9 --- .../webgpu/webgpu_provider_factory.cc | 42 +++++++++++-- .../providers/webgpu/webgpu_context_test.cc | 60 ++++++++++--------- 8 files changed, 88 insertions(+), 163 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index 86bbbbfc3a17f..b935a028a3ce6 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -42,16 +42,6 @@ class IExternalDataLoader { // writing into a tensor allocated by the framework. virtual bool CreatesTensorForDevice([[maybe_unused]] const OrtDevice& target_device) const { return false; } - // Determines whether a tensor-creating loader can use the allocator selected for - // the SessionState. An OK status with can_create_tensor=false lets ordinary - // initializer loading handle the tensor. - virtual common::Status CanCreateTensorWithAllocator( - [[maybe_unused]] const std::shared_ptr& allocator, - bool& can_create_tensor) const { - can_create_tensor = true; - return common::Status::OK(); - } - // Begins a batch covering the external initializers for one SessionState // initialization. ORT describes all loader-created tensors through PrepareTensor() // before calling FinalizeLoad(), allowing an implementation to coalesce file diff --git a/onnxruntime/core/framework/external_data_loader_manager.cc b/onnxruntime/core/framework/external_data_loader_manager.cc index 26ecf71e4254c..86746b4a2126c 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.cc +++ b/onnxruntime/core/framework/external_data_loader_manager.cc @@ -3,8 +3,6 @@ #include "core/framework/external_data_loader_manager.h" -#include "core/framework/allocator.h" - namespace onnxruntime { using namespace common; @@ -39,57 +37,16 @@ const IExternalDataLoader* ExternalDataLoaderManager::GetExternalDataLoader( } #if defined(ENABLE_D3D12_FILE_LOADING) -Status ExternalDataLoaderManager::GetTensorCreator( - const OrtDevice& target_device, - int32_t tensor_data_type, - const std::shared_ptr& allocator, - const IExternalDataLoader*& tensor_creator) const { - tensor_creator = nullptr; +const IExternalDataLoader* ExternalDataLoaderManager::GetTensorCreator( + const OrtDevice& target_device, int32_t tensor_data_type) const { for (const auto& external_data_loader : external_data_loaders_) { - if (!external_data_loader->SupportsDataType(tensor_data_type) || - !external_data_loader->CreatesTensorForDevice(target_device)) { - continue; - } - - bool can_create_tensor = false; - ORT_RETURN_IF_ERROR( - external_data_loader->CanCreateTensorWithAllocator(allocator, can_create_tensor)); - if (can_create_tensor) { - tensor_creator = external_data_loader.get(); - break; - } - } - - return Status::OK(); -} - -Status ExternalDataLoaderManager::GetExternalDataLoader( - const OrtMemoryInfo& target_memory_info, - int32_t tensor_data_type, - const std::shared_ptr& allocator, - const IExternalDataLoader*& selected_loader) const { - selected_loader = nullptr; - for (const auto& external_data_loader : external_data_loaders_) { - if (!external_data_loader->SupportsDataType(tensor_data_type)) { - continue; - } - - if (external_data_loader->CreatesTensorForDevice(target_memory_info.device)) { - bool can_create_tensor = false; - ORT_RETURN_IF_ERROR( - external_data_loader->CanCreateTensorWithAllocator(allocator, can_create_tensor)); - if (!can_create_tensor) { - continue; - } - } else if (!external_data_loader->CanLoad(target_memory_info)) { - continue; + if (external_data_loader->SupportsDataType(tensor_data_type) && + external_data_loader->CreatesTensorForDevice(target_device)) { + return external_data_loader.get(); } - - selected_loader = external_data_loader.get(); - break; } - return Status::OK(); + return nullptr; } Status ExternalDataLoaderManager::BeginLoad() const { diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index 5a5df43df8fdd..54c8abf5c5db0 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.h +++ b/onnxruntime/core/framework/external_data_loader_manager.h @@ -27,16 +27,8 @@ class ExternalDataLoaderManager { const IExternalDataLoader* GetExternalDataLoader( const OrtMemoryInfo& target_memory_info, int32_t tensor_data_type) const; #if defined(ENABLE_D3D12_FILE_LOADING) - common::Status GetTensorCreator( - const OrtDevice& target_device, - int32_t tensor_data_type, - const std::shared_ptr& allocator, - const IExternalDataLoader*& tensor_creator) const; - common::Status GetExternalDataLoader( - const OrtMemoryInfo& target_memory_info, - int32_t tensor_data_type, - const std::shared_ptr& allocator, - const IExternalDataLoader*& external_data_loader) const; + const IExternalDataLoader* GetTensorCreator( + const OrtDevice& target_device, int32_t tensor_data_type) const; common::Status BeginLoad() const; common::Status FinalizeLoad(const std::function& is_canceled) const; void AbortLoad() const noexcept; diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index e5dc3d1475d05..268dd1dd4286b 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -101,14 +101,8 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st if (utils::HasExternalData(tensor_proto) && !utils::HasExternalDataInMemory(tensor_proto)) { -#if defined(ENABLE_D3D12_FILE_LOADING) - const IExternalDataLoader* external_data_loader = nullptr; - ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetExternalDataLoader( - memory_info, tensor_proto.data_type(), alloc, external_data_loader)); -#else auto external_data_loader = external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); -#endif if (external_data_loader) { #if defined(ENABLE_D3D12_FILE_LOADING) if (external_data_loader->CreatesTensorForDevice(device)) { @@ -351,15 +345,11 @@ common::Status SaveInitializedTensors( // - we do not trace values that are in memory because they may be sitting on top of the user allocated // memory. #if defined(ENABLE_D3D12_FILE_LOADING) - const IExternalDataLoader* tensor_creator = nullptr; - if (utils::HasExternalData(*tensor_proto) && - !utils::HasExternalDataInMemory(*tensor_proto)) { - const auto& target_device = exec_plan.GetLocation(ort_value_index); - ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetTensorCreator( - target_device, tensor_proto->data_type(), - planner.GetInitializerAllocator(target_device), tensor_creator)); - } - const bool loader_creates_tensor = tensor_creator != nullptr; + const bool loader_creates_tensor = + utils::HasExternalData(*tensor_proto) && + !utils::HasExternalDataInMemory(*tensor_proto) && + external_data_loader_mgr.GetTensorCreator( + exec_plan.GetLocation(ort_value_index), tensor_proto->data_type()) != nullptr; #else constexpr bool loader_creates_tensor = false; #endif @@ -386,15 +376,10 @@ common::Status SaveInitializedTensors( } #if defined(ENABLE_D3D12_FILE_LOADING) if (utils::HasExternalData(*entry.second) && - !utils::HasExternalDataInMemory(*entry.second)) { - const auto& target_device = exec_plan.GetLocation(entry.first); - const IExternalDataLoader* tensor_creator = nullptr; - ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetTensorCreator( - target_device, entry.second->data_type(), - planner.GetInitializerAllocator(target_device), tensor_creator)); - if (tensor_creator != nullptr) { - continue; - } + !utils::HasExternalDataInMemory(*entry.second) && + external_data_loader_mgr.GetTensorCreator( + exec_plan.GetLocation(entry.first), entry.second->data_type()) != nullptr) { + continue; } #endif ORT_RETURN_IF_ERROR(planner.Trace(entry.first, entry.second)); @@ -443,11 +428,8 @@ common::Status SaveInitializedTensors( "Preparing session state weights is canceled due to user request."); } - const auto& target_device = exec_plan.GetLocation(entry.first); - const IExternalDataLoader* tensor_creator = nullptr; - ORT_RETURN_IF_ERROR(external_data_loader_mgr.GetTensorCreator( - target_device, entry.second->data_type(), - planner.GetInitializerAllocator(target_device), tensor_creator)); + const auto* tensor_creator = external_data_loader_mgr.GetTensorCreator( + exec_plan.GetLocation(entry.first), entry.second->data_type()); if (tensor_creator != nullptr) { ORT_RETURN_IF_ERROR( utils::PrepareExtDataForTensorFromTensorProto(env, graph_loc, *entry.second, *tensor_creator)); diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 2ff73c2f137e5..f8db5419dfe48 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -595,29 +595,6 @@ bool D3D12AcceleratedExternalDataLoader:: IsWeightLoadAccelerationRequired(impl_->mode)); } -common::Status ResolveWeightLoadAccelerationAllocator( - WeightLoadAccelerationMode mode, - const IAllocator* accelerated_allocator, - const std::shared_ptr& selected_allocator, - bool& compatible) { - compatible = selected_allocator != nullptr && - selected_allocator.get() == accelerated_allocator; - ORT_RETURN_IF( - !compatible && IsWeightLoadAccelerationRequired(mode), - "weightLoadAcceleration=\"required\" is incompatible with the selected " - "WebGPU initializer allocator."); - return common::Status::OK(); -} - -common::Status D3D12AcceleratedExternalDataLoader:: - CanCreateTensorWithAllocator( - const std::shared_ptr& allocator, - bool& can_create_tensor) const { - return ResolveWeightLoadAccelerationAllocator( - impl_->mode, impl_->state->impl_->allocator, - allocator, can_create_tensor); -} - common::Status D3D12AcceleratedExternalDataLoader::BeginLoad() const { AbortLoad(); diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index 7f4eadf7a10b3..7f10d63193f24 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -28,12 +28,6 @@ common::Status ResolveWeightLoadAccelerationMode( const common::Status& support_status, bool& enabled); -common::Status ResolveWeightLoadAccelerationAllocator( - WeightLoadAccelerationMode mode, - const IAllocator* accelerated_allocator, - const std::shared_ptr& selected_allocator, - bool& compatible); - // Shared by the loader and allocator so imported resources outlive initializer loading. class D3D12AcceleratedInitializerState { public: @@ -70,9 +64,6 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { bool CanLoad(const OrtMemoryInfo& target_memory_info) const override; bool SupportsDataType(int32_t tensor_data_type) const override; bool CreatesTensorForDevice(const OrtDevice& target_device) const override; - common::Status CanCreateTensorWithAllocator( - const std::shared_ptr& allocator, - bool& can_create_tensor) const override; common::Status BeginLoad() const override; common::Status PrepareTensor(const Env& env, const std::filesystem::path& data_file_path, diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 97742385bae5d..2d7434269cd14 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -78,7 +78,30 @@ WeightLoadAccelerationMode ParseWeightLoadAccelerationMode( ". Must be \"off\", \"preferred\", or \"required\"."); } -WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) { +WeightLoadAccelerationMode ResolveWeightLoadAccelerationMode( + const ConfigOptions& config_options) { + const auto mode = + ParseWeightLoadAccelerationMode(config_options); + const bool use_environment_allocators = + config_options.GetConfigOrDefault( + kOrtSessionOptionsConfigUseEnvAllocators, "0") == "1"; + if (!use_environment_allocators || + mode == WeightLoadAccelerationMode::Off) { + return mode; + } + ORT_ENFORCE( + mode != WeightLoadAccelerationMode::Required, + "weightLoadAcceleration=\"required\" is incompatible with " + "session.use_env_allocators=1."); + LOGS_DEFAULT(WARNING) + << "weightLoadAcceleration=\"preferred\" is disabled because " + "session.use_env_allocators=1."; + return WeightLoadAccelerationMode::Off; +} + +WebGpuExecutionProviderConfig ParseEpConfig( + const ConfigOptions& config_options, + WeightLoadAccelerationMode weight_load_acceleration_mode) { WebGpuExecutionProviderConfig webgpu_ep_config{}; if (std::string preferred_layout_str; @@ -104,7 +127,7 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) } webgpu_ep_config.weight_load_acceleration_mode = - ParseWeightLoadAccelerationMode(config_options); + weight_load_acceleration_mode; if (std::string pool_generations_str; config_options.TryGetConfigEntry(kSessionBufferPoolGenerations, pool_generations_str)) { @@ -212,11 +235,13 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) return webgpu_ep_config; } -WebGpuContextConfig ParseWebGpuContextConfig(const ConfigOptions& config_options) { +WebGpuContextConfig ParseWebGpuContextConfig( + const ConfigOptions& config_options, + WeightLoadAccelerationMode weight_load_acceleration_mode) { WebGpuContextConfig config{}; config.weight_load_acceleration_mode = - ParseWeightLoadAccelerationMode(config_options); + weight_load_acceleration_mode; if (std::string context_id_str; config_options.TryGetConfigEntry(kDeviceId, context_id_str)) { @@ -446,11 +471,16 @@ WebGpuContextConfig ParseWebGpuContextConfig(const ConfigOptions& config_options static std::shared_ptr CreateWebGpuProviderFactory( const ConfigOptions& config_options, uint64_t test_only_max_storage_buffer_binding_size) { + const auto weight_load_acceleration_mode = + ResolveWeightLoadAccelerationMode(config_options); + // prepare WebGpuExecutionProviderConfig - WebGpuExecutionProviderConfig webgpu_ep_config = ParseEpConfig(config_options); + WebGpuExecutionProviderConfig webgpu_ep_config = + ParseEpConfig(config_options, weight_load_acceleration_mode); // prepare WebGpuContextConfig - WebGpuContextConfig config = ParseWebGpuContextConfig(config_options); + WebGpuContextConfig config = + ParseWebGpuContextConfig(config_options, weight_load_acceleration_mode); config.test_only_max_storage_buffer_binding_size = test_only_max_storage_buffer_binding_size; // Load the Dawn library and create the WebGPU instance. diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 5fe9ebb7deac3..2e8e9ef524a7d 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -130,6 +130,13 @@ ConfigOptions CompileOnlyWeightLoadAccelerationOptions(const char* value) { return options; } +ConfigOptions EnvironmentAllocatorWeightLoadAccelerationOptions(const char* value) { + auto options = CompileOnlyWeightLoadAccelerationOptions(value); + ORT_THROW_IF_ERROR( + options.AddConfigEntry(kOrtSessionOptionsConfigUseEnvAllocators, "1")); + return options; +} + bool DeviceToggleIsEnabled(const webgpu::WebGpuContext& context, std::string_view toggle_name) { #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) const auto toggles = dawn::native::GetTogglesUsed(context.Device().Get()); @@ -1189,6 +1196,32 @@ TEST(WebGpuContextTest, RequiredWeightLoadAccelerationDoesNotExposeLoaderWithout true); } +TEST(WebGpuContextTest, PreferredWeightLoadAccelerationIsDisabledWithEnvironmentAllocators) { + auto ep = WebGpuProviderFactoryCreator::Create( + EnvironmentAllocatorWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Preferred)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); + EXPECT_EQ(ep->GetExternalDataLoader(), nullptr); +} + +TEST(WebGpuContextTest, RequiredWeightLoadAccelerationRejectsEnvironmentAllocators) { + EXPECT_THROW( + WebGpuProviderFactoryCreator::Create( + EnvironmentAllocatorWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Required)), + OnnxRuntimeException); +} + +TEST(WebGpuContextTest, DisabledWeightLoadAccelerationAllowsEnvironmentAllocators) { + auto ep = WebGpuProviderFactoryCreator::Create( + EnvironmentAllocatorWeightLoadAccelerationOptions( + kWeightLoadAcceleration_Off)) + ->CreateProvider(); + ASSERT_NE(ep, nullptr); + EXPECT_EQ(ep->GetExternalDataLoader(), nullptr); +} + #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) void ExpectD3D12BufferContents( ID3D12Device* device, @@ -1934,33 +1967,6 @@ TEST(WebGpuContextTest, WeightLoadAccelerationModeResolution) { EXPECT_EQ(required_status.ErrorMessage(), unsupported.ErrorMessage()); } -TEST(WebGpuContextTest, WeightLoadAccelerationAllocatorResolution) { - const auto accelerated_allocator = std::make_shared(); - const auto environment_allocator = std::make_shared(); - bool compatible = false; - - EXPECT_STATUS_OK(webgpu::ResolveWeightLoadAccelerationAllocator( - webgpu::WeightLoadAccelerationMode::Preferred, - accelerated_allocator.get(), accelerated_allocator, compatible)); - EXPECT_TRUE(compatible); - - compatible = true; - EXPECT_STATUS_OK(webgpu::ResolveWeightLoadAccelerationAllocator( - webgpu::WeightLoadAccelerationMode::Preferred, - accelerated_allocator.get(), environment_allocator, compatible)); - EXPECT_FALSE(compatible); - - compatible = true; - const auto required_status = - webgpu::ResolveWeightLoadAccelerationAllocator( - webgpu::WeightLoadAccelerationMode::Required, - accelerated_allocator.get(), environment_allocator, compatible); - EXPECT_FALSE(required_status.IsOK()); - EXPECT_FALSE(compatible); - EXPECT_NE(required_status.ErrorMessage().find( - "incompatible with the selected WebGPU initializer allocator"), - std::string::npos); -} #endif TEST(WebGpuContextTest, AdapterIndexRejectsConflictingSelectorOnReusedContext) { From 681b277baa09a11e17b16b6adb06860ead8cb5c8 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 18:12:30 -0700 Subject: [PATCH 30/46] Clarify external tensor loading contract Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/framework/external_data_loader.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index b935a028a3ce6..efdd96bf6dab7 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -97,7 +97,9 @@ class IExternalDataLoader { virtual void AbortLoad() const noexcept {} #endif - // A tensor-creating loader replaces tensor with one backed by memory owned through allocator. + // Tensor should be allocated with the correct memory info and size. A loader that + // creates tensors for the target device replaces it with one backed by memory + // owned through allocator. virtual common::Status LoadTensor([[maybe_unused]] const Env& env, [[maybe_unused]] const std::filesystem::path& data_file_path, [[maybe_unused]] std::string_view tensor_name, From 2deb4503f366a23c16120716534e518ddb121c73 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 19:01:52 -0700 Subject: [PATCH 31/46] Simplify accelerated initializer state Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../webgpu/d3d12_external_data_loader.cc | 50 ++++++++----------- .../webgpu/d3d12_external_data_loader.h | 20 +------- 2 files changed, 22 insertions(+), 48 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index f8db5419dfe48..e020fa6d8bca4 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -341,7 +341,11 @@ common::Status ResolveWeightLoadAccelerationMode( return common::Status::OK(); } -struct D3D12AcceleratedInitializerState::Impl { +// Shared by the loader and allocator so imported resources outlive initializer loading. +class D3D12AcceleratedInitializerState { + public: + ~D3D12AcceleratedInitializerState(); + std::mutex mutex; IAllocator* allocator = nullptr; std::unordered_map< @@ -350,25 +354,18 @@ struct D3D12AcceleratedInitializerState::Impl { imported_allocations; }; -D3D12AcceleratedInitializerState:: - D3D12AcceleratedInitializerState() = default; - D3D12AcceleratedInitializerState:: ~D3D12AcceleratedInitializerState() { - if (!impl_) { - return; - } - std::vector> allocations; { - std::lock_guard lock{impl_->mutex}; + std::lock_guard lock{mutex}; allocations.reserve( - impl_->imported_allocations.size()); - for (auto& entry : impl_->imported_allocations) { + imported_allocations.size()); + for (auto& entry : imported_allocations) { allocations.push_back(std::move(entry.second)); } - impl_->imported_allocations.clear(); + imported_allocations.clear(); } for (auto& allocation : allocations) { EndAccessNoThrow(*allocation); @@ -424,14 +421,14 @@ class D3D12AcceleratedWebGpuAllocator final std::unique_ptr imported; { std::lock_guard lock{ - state_->impl_->mutex}; + state_->mutex}; const auto iterator = - state_->impl_->imported_allocations.find( + state_->imported_allocations.find( static_cast(p)); if (iterator != - state_->impl_->imported_allocations.end()) { + state_->imported_allocations.end()) { imported = std::move(iterator->second); - state_->impl_->imported_allocations.erase( + state_->imported_allocations.erase( iterator); } } @@ -482,16 +479,12 @@ AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( std::shared_ptr& out_state) { auto state = - std::shared_ptr( - new D3D12AcceleratedInitializerState()); - state->impl_ = - std::make_unique< - D3D12AcceleratedInitializerState::Impl>(); + std::make_shared(); auto allocator = std::make_shared< D3D12AcceleratedWebGpuAllocator>( context, std::move(recording_getter), state); - state->impl_->allocator = allocator.get(); + state->allocator = allocator.get(); out_state = state; return allocator; } @@ -561,8 +554,7 @@ D3D12AcceleratedExternalDataLoader:: : impl_{std::make_unique( context, std::move(state), mode)} { ORT_ENFORCE( - impl_->state != nullptr && - impl_->state->impl_ != nullptr, + impl_->state != nullptr, "D3D12 accelerated allocator state is required."); } @@ -847,7 +839,7 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( "D3D12 accelerated initializer requires its device allocator."); ORT_RETURN_IF_NOT( allocator.get() == - impl_->state->impl_->allocator, + impl_->state->allocator, "D3D12 accelerated initializer was passed a different allocator."); const size_t length = @@ -915,14 +907,14 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( { std::lock_guard lock{ - impl_->state->impl_->mutex}; + impl_->state->mutex}; ORT_RETURN_IF( - impl_->state->impl_->imported_allocations + impl_->state->imported_allocations .find(buffer) != - impl_->state->impl_->imported_allocations + impl_->state->imported_allocations .end(), "D3D12 accelerated imported buffer was registered twice."); - impl_->state->impl_->imported_allocations.emplace( + impl_->state->imported_allocations.emplace( buffer, std::move(imported)); } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index 7f10d63193f24..bdacf05d4e8e5 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -28,25 +28,7 @@ common::Status ResolveWeightLoadAccelerationMode( const common::Status& support_status, bool& enabled); -// Shared by the loader and allocator so imported resources outlive initializer loading. -class D3D12AcceleratedInitializerState { - public: - ~D3D12AcceleratedInitializerState(); - - private: - struct Impl; - - D3D12AcceleratedInitializerState(); - - std::unique_ptr impl_; - - friend AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( - WebGpuContext& context, - std::function recording_getter, - std::shared_ptr& out_state); - friend class D3D12AcceleratedExternalDataLoader; - friend class D3D12AcceleratedWebGpuAllocator; -}; +class D3D12AcceleratedInitializerState; AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( WebGpuContext& context, From 39ed974eefba53af0dcfdd11237a4c30779ab94f Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 19:03:41 -0700 Subject: [PATCH 32/46] Clarify D3D12 loader event names Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../d3d12_file_buffer_loader.cc | 29 ++++++++++--------- .../d3d12_file_buffer_loader.h | 2 +- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index 00c7288ad311c..f7163721d955f 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -121,7 +121,7 @@ struct D3D12FileBufferLoader::UploadSlot { ComPtr resource; ComPtr allocator; ComPtr command_list; - wil::unique_handle read_event; + wil::unique_handle file_read_complete_event; void* mapped = nullptr; OVERLAPPED overlapped{}; HANDLE read_file = INVALID_HANDLE_VALUE; @@ -148,12 +148,12 @@ D3D12FileBufferLoader::~D3D12FileBufferLoader() { device_.Detach(); copy_queue_.Detach(); copy_fence_.Detach(); - (void)fence_event_.release(); + (void)copy_fence_complete_event_.release(); for (auto& slot : slots_) { slot.resource.Detach(); slot.allocator.Detach(); slot.command_list.Detach(); - (void)slot.read_event.release(); + (void)slot.file_read_complete_event.release(); } for (auto& heap : untracked_batch_.heaps) { heap.Detach(); @@ -213,9 +213,9 @@ common::Status D3D12FileBufferLoader::Initialize() { return HResultError("ID3D12Device::CreateFence", hr); } - fence_event_.reset(CreateEventExW( + copy_fence_complete_event_.reset(CreateEventExW( nullptr, nullptr, 0, EVENT_MODIFY_STATE | SYNCHRONIZE)); - if (!fence_event_) { + if (!copy_fence_complete_event_) { return Win32Error("CreateEventExW(copy fence)", GetLastError()); } @@ -257,10 +257,10 @@ common::Status D3D12FileBufferLoader::Initialize() { "ID3D12GraphicsCommandList::Close(initial)", hr); } - slot.read_event.reset(CreateEventExW( + slot.file_read_complete_event.reset(CreateEventExW( nullptr, nullptr, CREATE_EVENT_MANUAL_RESET, EVENT_MODIFY_STATE | SYNCHRONIZE)); - if (!slot.read_event) { + if (!slot.file_read_complete_event) { return Win32Error("CreateEventExW(read)", GetLastError()); } } @@ -656,14 +656,15 @@ common::Status D3D12FileBufferLoader::IssueRead( UploadSlot& slot, uint64_t offset, DWORD size) { - if (!ResetEvent(slot.read_event.get())) { + if (!ResetEvent(slot.file_read_complete_event.get())) { return Win32Error("ResetEvent(read)", GetLastError()); } slot.overlapped = {}; slot.overlapped.Offset = static_cast(offset); slot.overlapped.OffsetHigh = static_cast(offset >> 32); - slot.overlapped.hEvent = slot.read_event.get(); + slot.overlapped.hEvent = + slot.file_read_complete_event.get(); slot.file_offset = offset; slot.requested = size; slot.read_file = file; @@ -766,7 +767,7 @@ common::Status D3D12FileBufferLoader::ReadRegion( index < slots_.size(); ++index) { if (slots_[index].read_active) { handles[handle_count] = - slots_[index].read_event.get(); + slots_[index].file_read_complete_event.get(); slot_indices[handle_count] = index; ++handle_count; } @@ -1038,7 +1039,7 @@ common::Status D3D12FileBufferLoader::WaitForFence( } const HRESULT hr = copy_fence_->SetEventOnCompletion( - value, fence_event_.get()); + value, copy_fence_complete_event_.get()); if (FAILED(hr)) { return HResultError( "ID3D12Fence::SetEventOnCompletion", hr); @@ -1048,7 +1049,7 @@ common::Status D3D12FileBufferLoader::WaitForFence( kInitialCancellationPollMilliseconds; while (true) { const DWORD wait_result = WaitForSingleObject( - fence_event_.get(), + copy_fence_complete_event_.get(), cancellation_poll_milliseconds); if (wait_result != WAIT_OBJECT_0 && wait_result != WAIT_TIMEOUT) { @@ -1083,9 +1084,9 @@ void D3D12FileBufferLoader::WaitForFenceUncancelled( return; } if (SUCCEEDED(copy_fence_->SetEventOnCompletion( - value, fence_event_.get())) && + value, copy_fence_complete_event_.get())) && WaitForSingleObject( - fence_event_.get(), INFINITE) == + copy_fence_complete_event_.get(), INFINITE) == WAIT_OBJECT_0) { return; } diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h index 25efdb12e975d..03135effa9e21 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h @@ -144,7 +144,7 @@ class D3D12FileBufferLoader { Config config_; Microsoft::WRL::ComPtr copy_queue_; Microsoft::WRL::ComPtr copy_fence_; - wil::unique_handle fence_event_; + wil::unique_handle copy_fence_complete_event_; std::vector slots_; Batch untracked_batch_; uint64_t next_fence_value_ = 0; From e77cebdd5b61e3bc4882937a5086b35c8225250d Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Mon, 5 Oct 2026 19:17:00 -0700 Subject: [PATCH 33/46] Integrate D3D12 loader with current external data API Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/framework/external_data_loader.h | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index efdd96bf6dab7..123c05aa2e9b7 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -97,6 +97,7 @@ class IExternalDataLoader { virtual void AbortLoad() const noexcept {} #endif +#if defined(ENABLE_D3D12_FILE_LOADING) // Tensor should be allocated with the correct memory info and size. A loader that // creates tensors for the target device replaces it with one backed by memory // owned through allocator. From 647c8c8b92ba38282c557d6b8b5a17014bef82d6 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 09:34:45 -0700 Subject: [PATCH 34/46] Fix unused allocator in external data loading Mark the allocator parameter unused when D3D12 file loading is disabled so warning-as-error builds compile cleanly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/framework/tensorprotoutils.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 723610d984629..6118622c6325c 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -1960,6 +1960,10 @@ Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem: const IExternalDataLoader& ext_data_loader, const AllocatorPtr& allocator, Tensor& tensor) { +#if !defined(ENABLE_D3D12_FILE_LOADING) + ORT_UNUSED_PARAMETER(allocator); +#endif + ORT_ENFORCE(HasExternalData(tensor_proto)); // Defense-in-depth path validation for callers reaching this function outside Graph::Resolve. // In-memory markers are passed through; rejected explicitly below as unsupported for this path. From 2e6ac760d6255417b7789cc1203eaf691ba8e9ba Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:20:05 -0700 Subject: [PATCH 35/46] Clarify D3D12 file loader resource and operation names Use explicit source, read-plan, staging, and destination names. Store destination resources directly in BufferCollection. Detach destination resources when GPU completion cannot be tracked, and reject further loads. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../d3d12_file_buffer_loader.cc | 350 +++++++++--------- .../d3d12_file_buffer_loader.h | 74 ++-- .../webgpu/d3d12_external_data_loader.cc | 10 +- .../providers/webgpu/webgpu_context_test.cc | 13 +- 4 files changed, 216 insertions(+), 231 deletions(-) diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc index f7163721d955f..bc3aa37ceeb16 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -107,17 +107,17 @@ struct D3D12FileBufferLoader::FileReadRegion { uint64_t end = 0; }; -struct D3D12FileBufferLoader::PreparedFile { +struct D3D12FileBufferLoader::FileReadPlan { std::wstring path; wil::unique_hfile file; wil::unique_hfile buffered_tail_file; uint64_t size = 0; uint64_t alignment = 0; - std::vector range_indices; - std::vector regions; + std::vector buffer_source_indices; + std::vector read_regions; }; -struct D3D12FileBufferLoader::UploadSlot { +struct D3D12FileBufferLoader::UploadStagingSlot { ComPtr resource; ComPtr allocator; ComPtr command_list; @@ -126,14 +126,14 @@ struct D3D12FileBufferLoader::UploadSlot { OVERLAPPED overlapped{}; HANDLE read_file = INVALID_HANDLE_VALUE; uint64_t file_offset = 0; - DWORD requested = 0; + DWORD requested_read_bytes = 0; uint64_t fence_value = 0; bool read_active = false; }; struct D3D12FileBufferLoader::AllocationResult { common::Status status; - Batch batch; + BufferCollection batch; }; D3D12FileBufferLoader::D3D12FileBufferLoader( @@ -142,32 +142,26 @@ D3D12FileBufferLoader::D3D12FileBufferLoader( } D3D12FileBufferLoader::~D3D12FileBufferLoader() { - if (retain_untracked_submission_) { + if (upload_error_gpu_completion_unknown_) { // Submitted work has no usable completion fence, so its resources must // remain alive for the lifetime of the process. device_.Detach(); copy_queue_.Detach(); copy_fence_.Detach(); (void)copy_fence_complete_event_.release(); - for (auto& slot : slots_) { - slot.resource.Detach(); - slot.allocator.Detach(); - slot.command_list.Detach(); - (void)slot.file_read_complete_event.release(); - } - for (auto& heap : untracked_batch_.heaps) { - heap.Detach(); - } - for (auto& buffer : untracked_batch_.buffers) { - buffer.resource.Detach(); + for (auto& staging_slot : upload_staging_slots_) { + staging_slot.resource.Detach(); + staging_slot.allocator.Detach(); + staging_slot.command_list.Detach(); + (void)staging_slot.file_read_complete_event.release(); } return; } - for (auto& slot : slots_) { - if (slot.resource && slot.mapped != nullptr) { - slot.resource->Unmap(0, nullptr); - slot.mapped = nullptr; + for (auto& staging_slot : upload_staging_slots_) { + if (staging_slot.resource && staging_slot.mapped != nullptr) { + staging_slot.resource->Unmap(0, nullptr); + staging_slot.mapped = nullptr; } } } @@ -222,45 +216,45 @@ common::Status D3D12FileBufferLoader::Initialize() { const auto upload_heap = HeapProperties(D3D12_HEAP_TYPE_UPLOAD); const auto upload_description = BufferDescription(config_.upload_slot_size, D3D12_RESOURCE_FLAG_NONE); - slots_.resize(config_.upload_slot_count); - for (auto& slot : slots_) { + upload_staging_slots_.resize(config_.upload_slot_count); + for (auto& staging_slot : upload_staging_slots_) { hr = device_->CreateCommittedResource( &upload_heap, D3D12_HEAP_FLAG_NONE, &upload_description, D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, - IID_PPV_ARGS(&slot.resource)); + IID_PPV_ARGS(&staging_slot.resource)); if (FAILED(hr)) { return HResultError( "ID3D12Device::CreateCommittedResource(upload)", hr); } - hr = slot.resource->Map(0, nullptr, &slot.mapped); + hr = staging_slot.resource->Map(0, nullptr, &staging_slot.mapped); if (FAILED(hr)) { return HResultError("ID3D12Resource::Map(upload)", hr); } hr = device_->CreateCommandAllocator( - D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&slot.allocator)); + D3D12_COMMAND_LIST_TYPE_COPY, IID_PPV_ARGS(&staging_slot.allocator)); if (FAILED(hr)) { return HResultError( "ID3D12Device::CreateCommandAllocator", hr); } hr = device_->CreateCommandList( - 0, D3D12_COMMAND_LIST_TYPE_COPY, slot.allocator.Get(), nullptr, - IID_PPV_ARGS(&slot.command_list)); + 0, D3D12_COMMAND_LIST_TYPE_COPY, staging_slot.allocator.Get(), nullptr, + IID_PPV_ARGS(&staging_slot.command_list)); if (FAILED(hr)) { return HResultError("ID3D12Device::CreateCommandList", hr); } - hr = slot.command_list->Close(); + hr = staging_slot.command_list->Close(); if (FAILED(hr)) { return HResultError( "ID3D12GraphicsCommandList::Close(initial)", hr); } - slot.file_read_complete_event.reset(CreateEventExW( + staging_slot.file_read_complete_event.reset(CreateEventExW( nullptr, nullptr, CREATE_EVENT_MANUAL_RESET, EVENT_MODIFY_STATE | SYNCHRONIZE)); - if (!slot.file_read_complete_event) { + if (!staging_slot.file_read_complete_event) { return Win32Error("CreateEventExW(read)", GetLastError()); } } @@ -269,15 +263,15 @@ common::Status D3D12FileBufferLoader::Initialize() { } common::Status D3D12FileBufferLoader::LoadInternal( - const std::vector& ranges, - Batch& result, + const std::vector& ranges, + BufferCollection& result, const CancellationToken& cancellation) { std::lock_guard lock(load_mutex_); result.Clear(); ORT_RETURN_IF( - retain_untracked_submission_, - "D3D12 file buffer loader cannot be reused after an untracked command submission."); + upload_error_gpu_completion_unknown_, + "D3D12 file buffer loader cannot be reused when GPU completion is unknown after an upload error."); if (ranges.empty()) { return ORT_MAKE_STATUS( ONNXRUNTIME, INVALID_ARGUMENT, @@ -317,7 +311,7 @@ common::Status D3D12FileBufferLoader::LoadInternal( AllocationResult allocation; try { allocation.status = - AllocateDestinations(sizes, allocation.batch); + AllocateDestinationGpuBuffers(sizes, allocation.batch); } catch (const std::exception& ex) { allocation.status = ORT_MAKE_STATUS( ONNXRUNTIME, RUNTIME_EXCEPTION, @@ -330,7 +324,9 @@ common::Status D3D12FileBufferLoader::LoadInternal( return allocation; }); - std::vector files; + // CreateFileReadPlan groups buffer sources by file and coalesces overlapping + // or adjacent sector-aligned ranges into read regions to avoid redundant reads. + std::vector files; uint64_t last_submitted_fence = 0; bool load_succeeded = false; // Resources referenced by in-flight I/O or GPU work must remain alive on @@ -348,15 +344,15 @@ common::Status D3D12FileBufferLoader::LoadInternal( } }); - ORT_RETURN_IF_ERROR(PrepareFiles(ranges, cancellation, files)); + ORT_RETURN_IF_ERROR(CreateFileReadPlan(ranges, cancellation, files)); uint64_t maximum_alignment = 1; for (const auto& file : files) { maximum_alignment = std::max(maximum_alignment, file.alignment); } - ORT_RETURN_IF_ERROR(PrepareSlots(maximum_alignment, cancellation)); + ORT_RETURN_IF_ERROR(PrepareUploadStagingSlots(maximum_alignment, cancellation)); - Batch loaded_batch; + BufferCollection loaded_batch; bool allocation_ready = false; auto ensure_allocation = [&]() -> common::Status { if (allocation_ready) { @@ -372,8 +368,8 @@ common::Status D3D12FileBufferLoader::LoadInternal( }; for (auto& file : files) { - for (const auto& region : file.regions) { - const auto status = ReadRegion( + for (const auto& region : file.read_regions) { + const auto status = ReadFileRegion( file, region, ranges, loaded_batch, allocation_ready, ensure_allocation, cancellation, last_submitted_fence); if (!status.IsOK()) { @@ -386,16 +382,16 @@ common::Status D3D12FileBufferLoader::LoadInternal( ORT_RETURN_IF_ERROR( WaitForFence(last_submitted_fence, cancellation)); ORT_RETURN_IF_ERROR( - TransitionToCommon(loaded_batch, cancellation, last_submitted_fence)); + TransitionDestinationBuffersToCommonState(loaded_batch, cancellation, last_submitted_fence)); result = std::move(loaded_batch); load_succeeded = true; return common::Status::OK(); } -common::Status D3D12FileBufferLoader::AllocateDestinations( +common::Status D3D12FileBufferLoader::AllocateDestinationGpuBuffers( const std::vector& sizes, - Batch& batch) { + BufferCollection& batch) { struct Placement { bool committed = false; size_t heap_index = 0; @@ -486,13 +482,13 @@ common::Status D3D12FileBufferLoader::AllocateDestinations( hr = device_->CreateCommittedResource( &heap_properties, D3D12_HEAP_FLAG_NONE, &placement.description, D3D12_RESOURCE_STATE_COPY_DEST, nullptr, - IID_PPV_ARGS(&batch.buffers[index].resource)); + IID_PPV_ARGS(&batch.buffers[index])); } else { hr = device_->CreatePlacedResource( batch.heaps[placement.heap_index].Get(), placement.offset, &placement.description, D3D12_RESOURCE_STATE_COPY_DEST, nullptr, - IID_PPV_ARGS(&batch.buffers[index].resource)); + IID_PPV_ARGS(&batch.buffers[index])); } if (FAILED(hr)) { return HResultError( @@ -501,16 +497,15 @@ common::Status D3D12FileBufferLoader::AllocateDestinations( : "ID3D12Device::CreatePlacedResource", hr); } - batch.buffers[index].size = sizes[index]; } return common::Status::OK(); } -common::Status D3D12FileBufferLoader::PrepareFiles( - const std::vector& ranges, +common::Status D3D12FileBufferLoader::CreateFileReadPlan( + const std::vector& ranges, const CancellationToken& cancellation, - std::vector& files) { + std::vector& files) { // Group ranges by source file so each file is opened once. Later, aligned // overlapping ranges are merged to avoid redundant unbuffered reads. std::map file_indices; @@ -520,11 +515,11 @@ common::Status D3D12FileBufferLoader::PrepareFiles( const auto [file_index_it, inserted] = file_indices.try_emplace(range.path, files.size()); if (inserted) { - PreparedFile file; + FileReadPlan file; file.path = range.path; files.push_back(std::move(file)); } - files[file_index_it->second].range_indices.push_back(range_index); + files[file_index_it->second].buffer_source_indices.push_back(range_index); } for (auto& file : files) { @@ -589,12 +584,12 @@ common::Status D3D12FileBufferLoader::PrepareFiles( } std::sort( - file.range_indices.begin(), file.range_indices.end(), + file.buffer_source_indices.begin(), file.buffer_source_indices.end(), [&](size_t left, size_t right) { return ranges[left].offset < ranges[right].offset; }); - for (size_t range_index : file.range_indices) { + for (size_t range_index : file.buffer_source_indices) { const auto& range = ranges[range_index]; if (range.offset > file.size || range.length > file.size - range.offset) { @@ -615,12 +610,12 @@ common::Status D3D12FileBufferLoader::PrepareFiles( FileReadRegion region{ AlignDown(range.offset, file.alignment), std::min(aligned_end, file.size)}; - if (!file.regions.empty() && - region.begin <= file.regions.back().end) { - file.regions.back().end = - std::max(file.regions.back().end, region.end); + if (!file.read_regions.empty() && + region.begin <= file.read_regions.back().end) { + file.read_regions.back().end = + std::max(file.read_regions.back().end, region.end); } else { - file.regions.push_back(region); + file.read_regions.push_back(region); } } } @@ -628,76 +623,76 @@ common::Status D3D12FileBufferLoader::PrepareFiles( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::PrepareSlots( +common::Status D3D12FileBufferLoader::PrepareUploadStagingSlots( uint64_t alignment, const CancellationToken& cancellation) { - for (auto& slot : slots_) { + for (auto& staging_slot : upload_staging_slots_) { ORT_RETURN_IF_ERROR( - WaitForFence(slot.fence_value, cancellation)); - if (reinterpret_cast(slot.mapped) % + WaitForFence(staging_slot.fence_value, cancellation)); + if (reinterpret_cast(staging_slot.mapped) % alignment != 0) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, "Mapped upload buffer does not satisfy source file alignment."); } - slot.overlapped = {}; - slot.read_file = INVALID_HANDLE_VALUE; - slot.file_offset = 0; - slot.requested = 0; - slot.fence_value = 0; - slot.read_active = false; + staging_slot.overlapped = {}; + staging_slot.read_file = INVALID_HANDLE_VALUE; + staging_slot.file_offset = 0; + staging_slot.requested_read_bytes = 0; + staging_slot.fence_value = 0; + staging_slot.read_active = false; } return common::Status::OK(); } -common::Status D3D12FileBufferLoader::IssueRead( +common::Status D3D12FileBufferLoader::IssueFileRead( HANDLE file, - UploadSlot& slot, + UploadStagingSlot& staging_slot, uint64_t offset, DWORD size) { - if (!ResetEvent(slot.file_read_complete_event.get())) { + if (!ResetEvent(staging_slot.file_read_complete_event.get())) { return Win32Error("ResetEvent(read)", GetLastError()); } - slot.overlapped = {}; - slot.overlapped.Offset = static_cast(offset); - slot.overlapped.OffsetHigh = + staging_slot.overlapped = {}; + staging_slot.overlapped.Offset = static_cast(offset); + staging_slot.overlapped.OffsetHigh = static_cast(offset >> 32); - slot.overlapped.hEvent = - slot.file_read_complete_event.get(); - slot.file_offset = offset; - slot.requested = size; - slot.read_file = file; - slot.read_active = true; + staging_slot.overlapped.hEvent = + staging_slot.file_read_complete_event.get(); + staging_slot.file_offset = offset; + staging_slot.requested_read_bytes = size; + staging_slot.read_file = file; + staging_slot.read_active = true; if (!ReadFile( - file, slot.mapped, size, nullptr, - &slot.overlapped)) { + file, staging_slot.mapped, size, nullptr, + &staging_slot.overlapped)) { const DWORD error = GetLastError(); if (error != ERROR_IO_PENDING) { - slot.read_active = false; - slot.read_file = INVALID_HANDLE_VALUE; + staging_slot.read_active = false; + staging_slot.read_file = INVALID_HANDLE_VALUE; return Win32Error("ReadFile", error); } } return common::Status::OK(); } -common::Status D3D12FileBufferLoader::CompleteRead( +common::Status D3D12FileBufferLoader::ProcessFileReadCompletion( uint64_t file_size, - UploadSlot& slot, + UploadStagingSlot& staging_slot, DWORD& bytes_read) { if (!GetOverlappedResult( - slot.read_file, &slot.overlapped, + staging_slot.read_file, &staging_slot.overlapped, &bytes_read, FALSE)) { return Win32Error( "GetOverlappedResult", GetLastError()); } - slot.read_active = false; - slot.read_file = INVALID_HANDLE_VALUE; + staging_slot.read_active = false; + staging_slot.read_file = INVALID_HANDLE_VALUE; const uint64_t required_bytes = std::min( - slot.requested, - file_size - slot.file_offset); + staging_slot.requested_read_bytes, + file_size - staging_slot.file_offset); if (bytes_read < required_bytes) { return ORT_MAKE_STATUS( ONNXRUNTIME, FAIL, @@ -707,31 +702,31 @@ common::Status D3D12FileBufferLoader::CompleteRead( } template -common::Status D3D12FileBufferLoader::ReadRegion( - PreparedFile& file, +common::Status D3D12FileBufferLoader::ReadFileRegion( + FileReadPlan& file, const FileReadRegion& region, - const std::vector& ranges, - Batch& batch, + const std::vector& ranges, + BufferCollection& batch, bool& allocation_ready, EnsureAllocationFn& ensure_allocation, const CancellationToken& cancellation, uint64_t& last_submitted_fence) { uint64_t next_offset = region.begin; - // Each upload slot moves through read-active, copy-in-flight, and reusable - // states. Keep all eligible slots busy without reusing one before its - // previous GPU copy reaches the fence. + // Each upload staging slot moves through read-active, copy-in-flight, and + // reusable states. Keep all eligible staging slots busy without reusing one + // before its previous GPU copy reaches the fence. auto issue_available = [&]() -> common::Status { - for (auto& slot : slots_) { + for (auto& staging_slot : upload_staging_slots_) { if (next_offset >= region.end) { break; } - if (slot.read_active) { + if (staging_slot.read_active) { continue; } - if (slot.fence_value != 0 && + if (staging_slot.fence_value != 0 && copy_fence_->GetCompletedValue() < - slot.fence_value) { + staging_slot.fence_value) { continue; } if (cancellation.IsCancellationRequested()) { @@ -745,11 +740,11 @@ common::Status D3D12FileBufferLoader::ReadRegion( next_offset + bytes == file.size && bytes % file.alignment != 0; ORT_RETURN_IF_ERROR( - IssueRead( + IssueFileRead( use_buffered_tail ? file.buffered_tail_file.get() : file.file.get(), - slot, + staging_slot, next_offset, bytes)); next_offset += bytes; } @@ -761,14 +756,14 @@ common::Status D3D12FileBufferLoader::ReadRegion( kInitialCancellationPollMilliseconds; while (true) { std::array handles{}; - std::array slot_indices{}; + std::array staging_slot_indices{}; DWORD handle_count = 0; for (size_t index = 0; - index < slots_.size(); ++index) { - if (slots_[index].read_active) { + index < upload_staging_slots_.size(); ++index) { + if (upload_staging_slots_[index].read_active) { handles[handle_count] = - slots_[index].file_read_complete_event.get(); - slot_indices[handle_count] = index; + upload_staging_slots_[index].file_read_complete_event.get(); + staging_slot_indices[handle_count] = index; ++handle_count; } } @@ -779,11 +774,11 @@ common::Status D3D12FileBufferLoader::ReadRegion( } uint64_t next_fence = std::numeric_limits::max(); - for (const auto& slot : slots_) { - if (!slot.read_active && - slot.fence_value != 0) { + for (const auto& staging_slot : upload_staging_slots_) { + if (!staging_slot.read_active && + staging_slot.fence_value != 0) { next_fence = - std::min(next_fence, slot.fence_value); + std::min(next_fence, staging_slot.fence_value); } } if (next_fence == @@ -818,20 +813,20 @@ common::Status D3D12FileBufferLoader::ReadRegion( GetLastError()); } - auto& slot = - slots_[slot_indices[wait_result - WAIT_OBJECT_0]]; + auto& staging_slot = + upload_staging_slots_[staging_slot_indices[wait_result - WAIT_OBJECT_0]]; DWORD bytes_read = 0; ORT_RETURN_IF_ERROR( - CompleteRead( - file.size, slot, bytes_read)); + ProcessFileReadCompletion( + file.size, staging_slot, bytes_read)); if (!allocation_ready) { ORT_RETURN_IF_ERROR(ensure_allocation()); } ORT_RETURN_IF_ERROR( - SubmitCopies( - slot, slot.file_offset, bytes_read, - file.range_indices, ranges, batch, + SubmitStagingToDestinationCopies( + staging_slot, staging_slot.file_offset, bytes_read, + file.buffer_source_indices, ranges, batch, last_submitted_fence)); cancellation_poll_milliseconds = kInitialCancellationPollMilliseconds; @@ -841,28 +836,28 @@ common::Status D3D12FileBufferLoader::ReadRegion( return common::Status::OK(); } -common::Status D3D12FileBufferLoader::SubmitCopies( - UploadSlot& slot, +common::Status D3D12FileBufferLoader::SubmitStagingToDestinationCopies( + UploadStagingSlot& staging_slot, uint64_t chunk_begin, uint64_t bytes, - const std::vector& range_indices, - const std::vector& ranges, - Batch& batch, + const std::vector& buffer_source_indices, + const std::vector& ranges, + BufferCollection& batch, uint64_t& last_submitted_fence) { - HRESULT hr = slot.allocator->Reset(); + HRESULT hr = staging_slot.allocator->Reset(); if (FAILED(hr)) { return HResultError( "ID3D12CommandAllocator::Reset", hr); } - hr = slot.command_list->Reset( - slot.allocator.Get(), nullptr); + hr = staging_slot.command_list->Reset( + staging_slot.allocator.Get(), nullptr); if (FAILED(hr)) { return HResultError( "ID3D12GraphicsCommandList::Reset", hr); } const uint64_t chunk_end = chunk_begin + bytes; - for (size_t range_index : range_indices) { + for (size_t range_index : buffer_source_indices) { const auto& range = ranges[range_index]; const uint64_t range_end = range.offset + range.length; @@ -871,10 +866,10 @@ common::Status D3D12FileBufferLoader::SubmitCopies( const uint64_t copy_end = std::min(chunk_end, range_end); if (copy_begin < copy_end) { - slot.command_list->CopyBufferRegion( - batch.buffers[range_index].resource.Get(), + staging_slot.command_list->CopyBufferRegion( + batch.buffers[range_index].Get(), copy_begin - range.offset, - slot.resource.Get(), + staging_slot.resource.Get(), copy_begin - chunk_begin, copy_end - copy_begin); } @@ -895,10 +890,10 @@ common::Status D3D12FileBufferLoader::SubmitCopies( {aligned_size - padding_offset, range.length - source_relative_offset, chunk_end - source_offset}); - slot.command_list->CopyBufferRegion( - batch.buffers[range_index].resource.Get(), + staging_slot.command_list->CopyBufferRegion( + batch.buffers[range_index].Get(), padding_offset, - slot.resource.Get(), + staging_slot.resource.Get(), source_offset - chunk_begin, copy_size); padding_offset += copy_size; @@ -908,45 +903,38 @@ common::Status D3D12FileBufferLoader::SubmitCopies( } } - hr = slot.command_list->Close(); + hr = staging_slot.command_list->Close(); if (FAILED(hr)) { return HResultError( "ID3D12GraphicsCommandList::Close", hr); } ID3D12CommandList* command_lists[] = { - slot.command_list.Get()}; + staging_slot.command_list.Get()}; copy_queue_->ExecuteCommandLists(1, command_lists); const uint64_t fence_value = ++next_fence_value_; - const auto signal_status = - SignalSubmittedWork(fence_value); - if (!signal_status.IsOK()) { - if (retain_untracked_submission_) { - untracked_batch_ = std::move(batch); - } - return signal_status; - } - slot.fence_value = fence_value; + ORT_RETURN_IF_ERROR(SignalSubmittedWork(fence_value, batch)); + staging_slot.fence_value = fence_value; last_submitted_fence = fence_value; return common::Status::OK(); } -common::Status D3D12FileBufferLoader::TransitionToCommon( - Batch& batch, +common::Status D3D12FileBufferLoader::TransitionDestinationBuffersToCommonState( + BufferCollection& batch, const CancellationToken& cancellation, uint64_t& last_submitted_fence) { if (batch.buffers.empty()) { return common::Status::OK(); } - auto& slot = slots_.front(); - HRESULT hr = slot.allocator->Reset(); + auto& staging_slot = upload_staging_slots_.front(); + HRESULT hr = staging_slot.allocator->Reset(); if (FAILED(hr)) { return HResultError( "ID3D12CommandAllocator::Reset(transition)", hr); } - hr = slot.command_list->Reset( - slot.allocator.Get(), nullptr); + hr = staging_slot.command_list->Reset( + staging_slot.allocator.Get(), nullptr); if (FAILED(hr)) { return HResultError( "ID3D12GraphicsCommandList::Reset(transition)", hr); @@ -959,7 +947,7 @@ common::Status D3D12FileBufferLoader::TransitionToCommon( barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION; barrier.Transition.pResource = - buffer.resource.Get(); + buffer.Get(); barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES; barrier.Transition.StateBefore = @@ -968,40 +956,33 @@ common::Status D3D12FileBufferLoader::TransitionToCommon( D3D12_RESOURCE_STATE_COMMON; barriers.push_back(barrier); } - slot.command_list->ResourceBarrier( + staging_slot.command_list->ResourceBarrier( static_cast(barriers.size()), barriers.data()); - hr = slot.command_list->Close(); + hr = staging_slot.command_list->Close(); if (FAILED(hr)) { return HResultError( "ID3D12GraphicsCommandList::Close(transition)", hr); } ID3D12CommandList* command_lists[] = { - slot.command_list.Get()}; + staging_slot.command_list.Get()}; copy_queue_->ExecuteCommandLists(1, command_lists); const uint64_t fence_value = ++next_fence_value_; - const auto signal_status = - SignalSubmittedWork(fence_value); - if (!signal_status.IsOK()) { - if (retain_untracked_submission_) { - untracked_batch_ = std::move(batch); - } - return signal_status; - } + ORT_RETURN_IF_ERROR(SignalSubmittedWork(fence_value, batch)); last_submitted_fence = fence_value; return WaitForFence( last_submitted_fence, cancellation); } common::Status D3D12FileBufferLoader::SignalSubmittedWork( - uint64_t value) { + uint64_t value, BufferCollection& batch) { // ExecuteCommandLists has no status return. Once work has been submitted, // its resources must stay alive until either a fence is queued behind it or // the device is removed. If Signal fails while the device is still live, - // the submission cannot be tracked safely; retain the loader and submitted - // resources for process lifetime rather than hanging or releasing them. + // completion is unknown. Detach destination resources now; the destructor + // detaches the remaining GPU resources to keep them alive for process lifetime. const HRESULT signal_result = copy_queue_->Signal(copy_fence_.Get(), value); if (SUCCEEDED(signal_result)) { @@ -1014,7 +995,14 @@ common::Status D3D12FileBufferLoader::SignalSubmittedWork( "D3D12 device removed after command submission", removal_reason); } - retain_untracked_submission_ = true; + upload_error_gpu_completion_unknown_ = true; + for (auto& buffer : batch.buffers) { + buffer.Detach(); + } + for (auto& heap : batch.heaps) { + heap.Detach(); + } + batch.Clear(); return HResultError( "ID3D12CommandQueue::Signal after command submission", signal_result); @@ -1103,23 +1091,23 @@ void D3D12FileBufferLoader::WaitForFenceUncancelled( } void D3D12FileBufferLoader::DrainActiveReads() noexcept { - for (auto& slot : slots_) { - if (!slot.read_active) { + for (auto& staging_slot : upload_staging_slots_) { + if (!staging_slot.read_active) { continue; } (void)CancelIoEx( - slot.read_file, &slot.overlapped); + staging_slot.read_file, &staging_slot.overlapped); } - for (auto& slot : slots_) { - if (!slot.read_active) { + for (auto& staging_slot : upload_staging_slots_) { + if (!staging_slot.read_active) { continue; } DWORD bytes_read = 0; (void)GetOverlappedResult( - slot.read_file, &slot.overlapped, + staging_slot.read_file, &staging_slot.overlapped, &bytes_read, TRUE); - slot.read_active = false; - slot.read_file = INVALID_HANDLE_VALUE; + staging_slot.read_active = false; + staging_slot.read_file = INVALID_HANDLE_VALUE; } } @@ -1147,8 +1135,8 @@ common::Status D3D12FileBufferLoader::Create( } common::Status D3D12FileBufferLoader::Load( - const std::vector& ranges, - Batch& result, + const std::vector& ranges, + BufferCollection& result, const CancellationToken& cancellation) noexcept { try { return LoadInternal(ranges, result, cancellation); diff --git a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h index 03135effa9e21..961a5ca0c2bf8 100644 --- a/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h @@ -27,7 +27,7 @@ class D3D12FileBufferLoader { uint64_t max_destination_heap_size = 512ull * 1024ull * 1024ull; }; - struct FileRange { + struct BufferSource { std::wstring path; uint64_t offset = 0; uint64_t length = 0; @@ -45,14 +45,9 @@ class D3D12FileBufferLoader { } }; - struct Buffer { - Microsoft::WRL::ComPtr resource; - uint64_t size = 0; - }; - - struct Batch { + struct BufferCollection { std::vector> heaps; - std::vector buffers; + std::vector> buffers; void Clear() noexcept { buffers.clear(); @@ -70,17 +65,19 @@ class D3D12FileBufferLoader { D3D12FileBufferLoader(const D3D12FileBufferLoader&) = delete; D3D12FileBufferLoader& operator=(const D3D12FileBufferLoader&) = delete; + // Each BufferSource describes a file range for one returned buffer. // The returned buffers correspond one-to-one with ranges and preserve input order. // Each buffer is in D3D12_RESOURCE_STATE_COMMON when this method succeeds. + // Resource widths may include padding; ranges specify the loaded data lengths. common::Status Load( - const std::vector& ranges, - Batch& result, + const std::vector& ranges, + BufferCollection& result, const CancellationToken& cancellation = {}) noexcept; private: struct FileReadRegion; - struct PreparedFile; - struct UploadSlot; + struct FileReadPlan; + struct UploadStagingSlot; struct AllocationResult; D3D12FileBufferLoader( @@ -89,54 +86,54 @@ class D3D12FileBufferLoader { common::Status Initialize(); common::Status LoadInternal( - const std::vector& ranges, - Batch& result, + const std::vector& ranges, + BufferCollection& result, const CancellationToken& cancellation); - common::Status AllocateDestinations( + common::Status AllocateDestinationGpuBuffers( const std::vector& sizes, - Batch& batch); - common::Status PrepareFiles( - const std::vector& ranges, + BufferCollection& batch); + common::Status CreateFileReadPlan( + const std::vector& ranges, const CancellationToken& cancellation, - std::vector& files); - common::Status PrepareSlots( + std::vector& files); + common::Status PrepareUploadStagingSlots( uint64_t alignment, const CancellationToken& cancellation); - common::Status IssueRead( + common::Status IssueFileRead( HANDLE file, - UploadSlot& slot, + UploadStagingSlot& staging_slot, uint64_t offset, DWORD size); - common::Status CompleteRead( + common::Status ProcessFileReadCompletion( uint64_t file_size, - UploadSlot& slot, + UploadStagingSlot& staging_slot, DWORD& bytes_read); template - common::Status ReadRegion( - PreparedFile& file, + common::Status ReadFileRegion( + FileReadPlan& file, const FileReadRegion& region, - const std::vector& ranges, - Batch& batch, + const std::vector& ranges, + BufferCollection& batch, bool& allocation_ready, EnsureAllocationFn& ensure_allocation, const CancellationToken& cancellation, uint64_t& last_submitted_fence); - common::Status SubmitCopies( - UploadSlot& slot, + common::Status SubmitStagingToDestinationCopies( + UploadStagingSlot& staging_slot, uint64_t chunk_begin, uint64_t bytes, - const std::vector& range_indices, - const std::vector& ranges, - Batch& batch, + const std::vector& buffer_source_indices, + const std::vector& ranges, + BufferCollection& batch, uint64_t& last_submitted_fence); - common::Status TransitionToCommon( - Batch& batch, + common::Status TransitionDestinationBuffersToCommonState( + BufferCollection& batch, const CancellationToken& cancellation, uint64_t& last_submitted_fence); common::Status WaitForFence( uint64_t value, const CancellationToken& cancellation); - common::Status SignalSubmittedWork(uint64_t value); + common::Status SignalSubmittedWork(uint64_t value, BufferCollection& batch); void WaitForFenceUncancelled(uint64_t value) noexcept; void DrainActiveReads() noexcept; @@ -145,10 +142,9 @@ class D3D12FileBufferLoader { Microsoft::WRL::ComPtr copy_queue_; Microsoft::WRL::ComPtr copy_fence_; wil::unique_handle copy_fence_complete_event_; - std::vector slots_; - Batch untracked_batch_; + std::vector upload_staging_slots_; uint64_t next_fence_value_ = 0; - bool retain_untracked_submission_ = false; + bool upload_error_gpu_completion_unknown_ = false; std::mutex load_mutex_; }; diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index e020fa6d8bca4..1d1e68e3b1cf2 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -75,7 +75,7 @@ struct TensorKey { struct PreparedTensor { TensorKey key; ComPtr resource; - std::shared_ptr heap_owner; + std::shared_ptr heap_owner; wgpu::SharedBufferMemory memory; wgpu::Buffer buffer; bool access_started = false; @@ -144,7 +144,7 @@ common::Status LoadBatchToD3D12( "D3D12 accelerated initializer loading was canceled."); } - std::vector ranges; + std::vector ranges; std::vector tensor_indices; ranges.reserve(batch.load_range_count); tensor_indices.reserve(batch.load_range_count); @@ -167,7 +167,7 @@ common::Status LoadBatchToD3D12( return common::Status::OK(); } - auto loaded = std::make_shared(); + auto loaded = std::make_shared(); const CancellationState cancellation_state{ &is_canceled, &abort_requested}; const D3D12FileBufferLoader::CancellationToken cancellation{ @@ -183,7 +183,7 @@ common::Status LoadBatchToD3D12( for (size_t index = 0; index < tensor_indices.size(); ++index) { auto& tensor = batch.tensors[tensor_indices[index]]; - tensor.resource = loaded->buffers[index].resource; + tensor.resource = loaded->buffers[index]; tensor.heap_owner = loaded; } @@ -280,7 +280,7 @@ common::Status PrepareTensorForBatch( struct ImportedAllocation { ComPtr resource; - std::shared_ptr heap_owner; + std::shared_ptr heap_owner; wgpu::SharedBufferMemory memory; wgpu::Buffer buffer; bool access_started = false; diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 2e8e9ef524a7d..60216673b300e 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1225,13 +1225,14 @@ TEST(WebGpuContextTest, DisabledWeightLoadAccelerationAllowsEnvironmentAllocator #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) void ExpectD3D12BufferContents( ID3D12Device* device, - const windows::d3d12::D3D12FileBufferLoader::Batch& batch, + const windows::d3d12::D3D12FileBufferLoader::BufferCollection& batch, const std::vector>& expected_buffers) { ASSERT_EQ(batch.buffers.size(), expected_buffers.size()); uint64_t total_size = 0; for (size_t index = 0; index < expected_buffers.size(); ++index) { - ASSERT_LE(expected_buffers[index].size(), batch.buffers[index].size); + ASSERT_NE(batch.buffers[index].Get(), nullptr); + ASSERT_LE(expected_buffers[index].size(), batch.buffers[index]->GetDesc().Width); total_size += expected_buffers[index].size(); } ASSERT_GT(total_size, 0u); @@ -1270,7 +1271,7 @@ void ExpectD3D12BufferContents( for (size_t index = 0; index < expected_buffers.size(); ++index) { command_list->CopyBufferRegion( readback.Get(), destination_offset, - batch.buffers[index].resource.Get(), 0, + batch.buffers[index].Get(), 0, expected_buffers[index].size()); destination_offset += expected_buffers[index].size(); } @@ -1376,7 +1377,7 @@ TEST(WebGpuContextTest, D3D12FileLoaderReusesUploadSlotAcrossChunks) { ASSERT_STATUS_OK(windows::d3d12::D3D12FileBufferLoader::Create( context.WeightLoadingD3D12Device(), loader, config)); - windows::d3d12::D3D12FileBufferLoader::Batch batch; + windows::d3d12::D3D12FileBufferLoader::BufferCollection batch; ASSERT_STATUS_OK(loader->Load( {{data_path.native(), 0, kDataSize}}, batch)); ASSERT_EQ(batch.buffers.size(), 1u); @@ -1397,7 +1398,7 @@ TEST(WebGpuContextTest, D3D12FileLoaderUsesMultipleHeapsAndCommittedResources) { const auto data_path = std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); - std::vector ranges; + std::vector ranges; uint64_t file_offset = 0; { std::ofstream stream{data_path, std::ios::binary | std::ios::trunc}; @@ -1431,7 +1432,7 @@ TEST(WebGpuContextTest, D3D12FileLoaderUsesMultipleHeapsAndCommittedResources) { ASSERT_STATUS_OK(windows::d3d12::D3D12FileBufferLoader::Create( context.WeightLoadingD3D12Device(), loader, config)); - windows::d3d12::D3D12FileBufferLoader::Batch batch; + windows::d3d12::D3D12FileBufferLoader::BufferCollection batch; ASSERT_STATUS_OK(loader->Load(ranges, batch)); ASSERT_EQ(batch.buffers.size(), expected_buffers.size()); EXPECT_EQ(batch.heaps.size(), 2u); From 1f2a10536159a1cac5fca01c1286b9f4f635d401 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:28:22 -0700 Subject: [PATCH 36/46] Remove unnecessary preprocessor directive for D3D12 file loading in IExternalDataLoader --- onnxruntime/core/framework/external_data_loader.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index 123c05aa2e9b7..7fa17019c9ae2 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -95,9 +95,7 @@ class IExternalDataLoader { // PrepareTensor(), FinalizeLoad(), or partial tensor consumption, and may be called // more than once. Implementations must therefore be noexcept and idempotent. virtual void AbortLoad() const noexcept {} -#endif -#if defined(ENABLE_D3D12_FILE_LOADING) // Tensor should be allocated with the correct memory info and size. A loader that // creates tensors for the target device replaces it with one backed by memory // owned through allocator. From 17c181c93f9d52fa19a844e786bbfbd5218cbd1c Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:33:44 -0700 Subject: [PATCH 37/46] Expose external tensor loader overload in WebAssembly builds Make the allocator-aware LoadTensor overload available to WebAssembly loaders while restricting batch preparation hooks to D3D12 file loading. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/framework/external_data_loader.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index 7fa17019c9ae2..c690a477b23e8 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -95,10 +95,12 @@ class IExternalDataLoader { // PrepareTensor(), FinalizeLoad(), or partial tensor consumption, and may be called // more than once. Implementations must therefore be noexcept and idempotent. virtual void AbortLoad() const noexcept {} +#endif // Tensor should be allocated with the correct memory info and size. A loader that // creates tensors for the target device replaces it with one backed by memory // owned through allocator. +#if defined(ENABLE_D3D12_FILE_LOADING) || defined(__wasm__) virtual common::Status LoadTensor([[maybe_unused]] const Env& env, [[maybe_unused]] const std::filesystem::path& data_file_path, [[maybe_unused]] std::string_view tensor_name, From 60589c4d657dd4b954e9174e875d31aebc812e69 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 17:05:21 -0700 Subject: [PATCH 38/46] Document WebGPU copy recording synchronization Explain why copy encoding shares the recording lock with public allocator operations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/providers/webgpu/data_transfer.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 4b8ec7c523043..2ad2ac9621b28 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -17,6 +17,8 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data, bool dst_is_gpu, size_t bytes) const { auto& command_state = recording_; + // Public allocator calls can share this recorder; serialize copy encoding, + // not just submission, against allocation and release. std::lock_guard lock{command_state.mutex}; if (bytes > 0) { if (dst_is_gpu) { From 4c930226b153bc99cd5e0d58b97594481933e60d Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 17:17:47 -0700 Subject: [PATCH 39/46] Refactor D3D12 accelerated external data loader to remove Impl struct and streamline member variables --- .../webgpu/d3d12_external_data_loader.cc | 176 ++++++++---------- .../webgpu/d3d12_external_data_loader.h | 26 ++- 2 files changed, 96 insertions(+), 106 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 1d1e68e3b1cf2..daa3e7f5b332b 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -82,6 +82,8 @@ struct PreparedTensor { bool claimed = false; }; +} // namespace + struct D3D12AcceleratedBatch { struct FileInfo { std::filesystem::path canonical_path; @@ -96,6 +98,8 @@ struct D3D12AcceleratedBatch { bool finalized = false; }; +namespace { + struct D3D12AcceleratedLoadMetrics { double load_ms = 0.0; size_t loaded_ranges = 0; @@ -489,60 +493,36 @@ AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( return allocator; } -struct D3D12AcceleratedExternalDataLoader::Impl { - Impl( - WebGpuContext& context_in, - std::shared_ptr - state_in, - WeightLoadAccelerationMode mode_in) - : context{context_in}, - state{std::move(state_in)}, - mode{mode_in} { - } - - void ResolveSupport() const { - std::call_once(support_once, [this]() { - common::Status support_status = - common::Status::OK(); - if (IsWeightLoadAccelerationEnabled(mode)) { - support_status = - CheckD3D12AcceleratedExternalWeightsSupport( - context); - } - resolved_status = - ResolveWeightLoadAccelerationMode( - mode, support_status, enabled); - if (resolved_status.IsOK() && !enabled && - mode == WeightLoadAccelerationMode::Preferred) { - LOGS_DEFAULT(WARNING) - << "D3D12 accelerated external weights are unavailable; " - "using the ordinary WebGPU initializer loading path. " - "Reason: " - << support_status.ErrorMessage(); - } - }); - } - - common::Status EnsureFileLoader() const { - if (file_loader) { - return common::Status::OK(); +void D3D12AcceleratedExternalDataLoader::ResolveSupport() const { + std::call_once(support_once_, [this]() { + common::Status support_status = + common::Status::OK(); + if (IsWeightLoadAccelerationEnabled(mode_)) { + support_status = + CheckD3D12AcceleratedExternalWeightsSupport( + context_); } - return D3D12FileBufferLoader::Create( - context.WeightLoadingD3D12Device(), file_loader); - } + resolved_status_ = + ResolveWeightLoadAccelerationMode( + mode_, support_status, enabled_); + if (resolved_status_.IsOK() && !enabled_ && + mode_ == WeightLoadAccelerationMode::Preferred) { + LOGS_DEFAULT(WARNING) + << "D3D12 accelerated external weights are unavailable; " + "using the ordinary WebGPU initializer loading path. " + "Reason: " + << support_status.ErrorMessage(); + } + }); +} - WebGpuContext& context; - std::shared_ptr - state; - WeightLoadAccelerationMode mode; - mutable std::once_flag support_once; - mutable common::Status resolved_status; - mutable bool enabled = false; - mutable std::atomic abort_requested{false}; - mutable std::unique_ptr - file_loader; - mutable std::unique_ptr batch; -}; +common::Status D3D12AcceleratedExternalDataLoader::EnsureFileLoader() const { + if (file_loader_) { + return common::Status::OK(); + } + return D3D12FileBufferLoader::Create( + context_.WeightLoadingD3D12Device(), file_loader_); +} D3D12AcceleratedExternalDataLoader:: D3D12AcceleratedExternalDataLoader( @@ -551,10 +531,9 @@ D3D12AcceleratedExternalDataLoader:: D3D12AcceleratedInitializerState> state, WeightLoadAccelerationMode mode) - : impl_{std::make_unique( - context, std::move(state), mode)} { + : context_{context}, state_{std::move(state)}, mode_{mode} { ORT_ENFORCE( - impl_->state != nullptr, + state_ != nullptr, "D3D12 accelerated allocator state is required."); } @@ -565,10 +544,10 @@ D3D12AcceleratedExternalDataLoader:: bool D3D12AcceleratedExternalDataLoader::CanLoad( const OrtMemoryInfo& target_memory_info) const { - impl_->ResolveSupport(); + ResolveSupport(); return target_memory_info.device == WebGpuDevice && target_memory_info.name == WEBGPU_BUFFER && - impl_->enabled; + enabled_; } bool D3D12AcceleratedExternalDataLoader:: @@ -581,23 +560,23 @@ bool D3D12AcceleratedExternalDataLoader:: bool D3D12AcceleratedExternalDataLoader:: CreatesTensorForDevice( const OrtDevice& target_device) const { - impl_->ResolveSupport(); + ResolveSupport(); return target_device == WebGpuDevice && - (impl_->enabled || - IsWeightLoadAccelerationRequired(impl_->mode)); + (enabled_ || + IsWeightLoadAccelerationRequired(mode_)); } common::Status D3D12AcceleratedExternalDataLoader::BeginLoad() const { AbortLoad(); - impl_->ResolveSupport(); - ORT_RETURN_IF_ERROR(impl_->resolved_status); - if (!impl_->enabled) { + ResolveSupport(); + ORT_RETURN_IF_ERROR(resolved_status_); + if (!enabled_) { return common::Status::OK(); } - impl_->abort_requested.store( + abort_requested_.store( false, std::memory_order_relaxed); - impl_->batch = + batch_ = std::make_unique(); return common::Status::OK(); } @@ -610,25 +589,25 @@ D3D12AcceleratedExternalDataLoader::PrepareTensor( FileOffsetType data_offset, SafeInt data_length) const { ORT_RETURN_IF_NOT( - impl_->batch != nullptr && - !impl_->batch->finalized, + batch_ != nullptr && + !batch_->finalized, "D3D12 accelerated initializer batch has not been started."); return PrepareTensorForBatch( - *impl_->batch, env, data_file_path, + *batch_, env, data_file_path, tensor_name, data_offset, data_length); } common::Status D3D12AcceleratedExternalDataLoader::FinalizeLoad( const std::function& is_canceled) const { - if (!impl_->enabled) { + if (!enabled_) { return common::Status::OK(); } ORT_RETURN_IF_NOT( - impl_->batch != nullptr && - !impl_->batch->finalized, + batch_ != nullptr && + !batch_->finalized, "D3D12 accelerated initializer batch has not been started."); - auto& batch = *impl_->batch; + auto& batch = *batch_; const auto fail_or_fallback = [this, &is_canceled]( const common::Status& status) @@ -641,14 +620,14 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( "D3D12 accelerated initializer loading was canceled."); } if (IsWeightLoadAccelerationRequired( - impl_->mode)) { + mode_)) { return status; } LOGS_DEFAULT(WARNING) << "D3D12 accelerated initializer loading failed; " "using the ordinary WebGPU initializer path: " << status.ErrorMessage(); - impl_->enabled = false; + enabled_ = false; AbortLoad(); return common::Status::OK(); }; @@ -664,13 +643,13 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( } D3D12AcceleratedLoadMetrics load_metrics; - const auto loader_status = impl_->EnsureFileLoader(); + const auto loader_status = EnsureFileLoader(); if (!loader_status.IsOK()) { return fail_or_fallback(loader_status); } const auto load_status = LoadBatchToD3D12( - batch, *impl_->file_loader, is_canceled, - impl_->abort_requested, load_metrics); + batch, *file_loader_, is_canceled, + abort_requested_, load_metrics); if (!load_status.IsOK()) { return fail_or_fallback(load_status); } @@ -685,7 +664,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( const auto import_status = [&]() -> common::Status { ORT_RETURN_IF_NOT( - impl_->context.D3D12SharedResourceFeaturesAvailable(), + context_.D3D12SharedResourceFeaturesAvailable(), "D3D12 accelerated external weights require Dawn " "SharedBufferMemoryD3D12Resource and SharedFenceDXGISharedHandle features."); for (auto& tensor : batch.tensors) { @@ -708,7 +687,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( memory_descriptor.nextInChain = &resource_descriptor; tensor.memory = - impl_->context.Device() + context_.Device() .ImportSharedBufferMemory( &memory_descriptor); ORT_RETURN_IF_NOT( @@ -782,18 +761,13 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( void D3D12AcceleratedExternalDataLoader::AbortLoad() const noexcept { - if (impl_) { - impl_->abort_requested.store( - true, std::memory_order_relaxed); - } - if (!impl_) { - return; - } + abort_requested_.store( + true, std::memory_order_relaxed); ORT_TRY { - if (impl_->batch) { + if (batch_) { for (auto& tensor : - impl_->batch->tensors) { + batch_->tensors) { if (!tensor.claimed && tensor.access_started) { ImportedAllocation allocation; @@ -810,7 +784,7 @@ void D3D12AcceleratedExternalDataLoader::AbortLoad() tensor.access_started = false; } } - impl_->batch.reset(); + batch_.reset(); } } ORT_CATCH(...) { @@ -828,8 +802,8 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( const std::shared_ptr& allocator, Tensor& tensor) const { ORT_RETURN_IF_NOT( - impl_->batch != nullptr && - impl_->batch->finalized, + batch_ != nullptr && + batch_->finalized, "D3D12 accelerated initializer batch has not been finalized."); ORT_RETURN_IF( data_offset < 0, @@ -839,16 +813,16 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( "D3D12 accelerated initializer requires its device allocator."); ORT_RETURN_IF_NOT( allocator.get() == - impl_->state->allocator, + state_->allocator, "D3D12 accelerated initializer was passed a different allocator."); const size_t length = static_cast(data_length); const auto file_iterator = - impl_->batch->files.find( + batch_->files.find( data_file_path.lexically_normal()); ORT_RETURN_IF( - file_iterator == impl_->batch->files.end(), + file_iterator == batch_->files.end(), "No prepared D3D12 data file matches \"", data_file_path.string(), "\"."); const TensorKey key = MakeKey( @@ -857,16 +831,16 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( static_cast(data_offset), length); const auto iterator = - impl_->batch->tensors_by_key.find(key); + batch_->tensors_by_key.find(key); ORT_RETURN_IF( iterator == - impl_->batch->tensors_by_key.end(), + batch_->tensors_by_key.end(), "No prepared D3D12 accelerated initializer matches \"", tensor_name, "\" and the requested file range."); auto& prepared = - impl_->batch->tensors[iterator->second]; + batch_->tensors[iterator->second]; ORT_RETURN_IF( prepared.claimed, "D3D12 accelerated initializer \"", @@ -907,14 +881,14 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( { std::lock_guard lock{ - impl_->state->mutex}; + state_->mutex}; ORT_RETURN_IF( - impl_->state->imported_allocations + state_->imported_allocations .find(buffer) != - impl_->state->imported_allocations + state_->imported_allocations .end(), "D3D12 accelerated imported buffer was registered twice."); - impl_->state->imported_allocations.emplace( + state_->imported_allocations.emplace( buffer, std::move(imported)); } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index bdacf05d4e8e5..b17a9ba9e2f36 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -5,17 +5,22 @@ #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +#include #include #include -#include +#include #include "core/framework/allocator.h" #include "core/framework/external_data_loader.h" #include "core/providers/webgpu/webgpu_provider_options.h" -struct ID3D12Device; - namespace onnxruntime { +namespace windows { +namespace d3d12 { +class D3D12FileBufferLoader; +} // namespace d3d12 +} // namespace windows + namespace webgpu { class WebGpuContext; @@ -29,6 +34,7 @@ common::Status ResolveWeightLoadAccelerationMode( bool& enabled); class D3D12AcceleratedInitializerState; +struct D3D12AcceleratedBatch; AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( WebGpuContext& context, @@ -63,8 +69,18 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { Tensor& tensor) const override; private: - struct Impl; - std::unique_ptr impl_; + void ResolveSupport() const; + common::Status EnsureFileLoader() const; + + WebGpuContext& context_; + std::shared_ptr state_; + WeightLoadAccelerationMode mode_; + mutable std::once_flag support_once_; + mutable common::Status resolved_status_; + mutable bool enabled_ = false; + mutable std::atomic abort_requested_{false}; + mutable std::unique_ptr file_loader_; + mutable std::unique_ptr batch_; }; } // namespace webgpu From 0efb67433201d67589939ba1957cf902f4b3a36d Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Tue, 6 Oct 2026 17:48:56 -0700 Subject: [PATCH 40/46] Refactor D3D12 file loading structures and update related references for clarity and consistency --- .../webgpu/d3d12_external_data_loader.cc | 142 +++++++++--------- .../webgpu/d3d12_external_data_loader.h | 23 +-- .../webgpu/webgpu_execution_provider.h | 4 +- .../providers/webgpu/webgpu_context_test.cc | 79 +++++++--- 4 files changed, 143 insertions(+), 105 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index daa3e7f5b332b..87dcdf8229aae 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -72,7 +72,7 @@ struct TensorKey { } }; -struct PreparedTensor { +struct AcceleratedTensorLoadEntry { TensorKey key; ComPtr resource; std::shared_ptr heap_owner; @@ -84,13 +84,13 @@ struct PreparedTensor { } // namespace -struct D3D12AcceleratedBatch { +struct D3D12AcceleratedLoadBatch { struct FileInfo { std::filesystem::path canonical_path; size_t length; }; - std::vector tensors; + std::vector tensors; std::map tensors_by_key; std::map files; size_t total_bytes = 0; @@ -136,7 +136,7 @@ bool IsCancellationRequested(void* opaque) noexcept { } common::Status LoadBatchToD3D12( - D3D12AcceleratedBatch& batch, + D3D12AcceleratedLoadBatch& batch, D3D12FileBufferLoader& loader, const std::function& is_canceled, const std::atomic& abort_requested, @@ -204,7 +204,7 @@ TensorKey MakeKey(const std::filesystem::path& path, } common::Status PrepareTensorForBatch( - D3D12AcceleratedBatch& batch, + D3D12AcceleratedLoadBatch& batch, const Env& env, const std::filesystem::path& data_file_path, std::string_view tensor_name, @@ -242,7 +242,7 @@ common::Status PrepareTensorForBatch( batch.files .emplace( file_key, - D3D12AcceleratedBatch::FileInfo{ + D3D12AcceleratedLoadBatch::FileInfo{ normalized_path, file_length}) .first; } @@ -346,9 +346,9 @@ common::Status ResolveWeightLoadAccelerationMode( } // Shared by the loader and allocator so imported resources outlive initializer loading. -class D3D12AcceleratedInitializerState { +class D3D12ImportedBufferRegistry { public: - ~D3D12AcceleratedInitializerState(); + ~D3D12ImportedBufferRegistry(); std::mutex mutex; IAllocator* allocator = nullptr; @@ -358,8 +358,8 @@ class D3D12AcceleratedInitializerState { imported_allocations; }; -D3D12AcceleratedInitializerState:: - ~D3D12AcceleratedInitializerState() { +D3D12ImportedBufferRegistry:: + ~D3D12ImportedBufferRegistry() { std::vector> allocations; { @@ -383,8 +383,8 @@ class D3D12AcceleratedWebGpuAllocator final WebGpuContext& context, std::function recording_getter, - std::shared_ptr - state) + std::shared_ptr + buffer_registry) : IAllocator( OrtMemoryInfo( WEBGPU_BUFFER, @@ -394,7 +394,7 @@ class D3D12AcceleratedWebGpuAllocator final context_{context}, recording_getter_{ std::move(recording_getter)}, - state_{std::move(state)} { + buffer_registry_{std::move(buffer_registry)} { } void* Alloc(size_t size) override { @@ -425,14 +425,14 @@ class D3D12AcceleratedWebGpuAllocator final std::unique_ptr imported; { std::lock_guard lock{ - state_->mutex}; + buffer_registry_->mutex}; const auto iterator = - state_->imported_allocations.find( + buffer_registry_->imported_allocations.find( static_cast(p)); if (iterator != - state_->imported_allocations.end()) { + buffer_registry_->imported_allocations.end()) { imported = std::move(iterator->second); - state_->imported_allocations.erase( + buffer_registry_->imported_allocations.erase( iterator); } } @@ -472,41 +472,41 @@ class D3D12AcceleratedWebGpuAllocator final WebGpuContext& context_; std::function recording_getter_; - std::shared_ptr - state_; + std::shared_ptr + buffer_registry_; }; AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( WebGpuContext& context, std::function recording_getter, - std::shared_ptr& - out_state) { - auto state = - std::make_shared(); + std::shared_ptr& + out_buffer_registry) { + auto buffer_registry = + std::make_shared(); auto allocator = std::make_shared< D3D12AcceleratedWebGpuAllocator>( - context, std::move(recording_getter), state); - state->allocator = allocator.get(); - out_state = state; + context, std::move(recording_getter), buffer_registry); + buffer_registry->allocator = allocator.get(); + out_buffer_registry = buffer_registry; return allocator; } void D3D12AcceleratedExternalDataLoader::ResolveSupport() const { - std::call_once(support_once_, [this]() { + std::call_once(device_support_resolution_once_, [this]() { common::Status support_status = common::Status::OK(); - if (IsWeightLoadAccelerationEnabled(mode_)) { + if (IsWeightLoadAccelerationEnabled(acceleration_mode_)) { support_status = CheckD3D12AcceleratedExternalWeightsSupport( context_); } - resolved_status_ = + device_support_resolution_status_ = ResolveWeightLoadAccelerationMode( - mode_, support_status, enabled_); - if (resolved_status_.IsOK() && !enabled_ && - mode_ == WeightLoadAccelerationMode::Preferred) { + acceleration_mode_, support_status, resolved_acceleration_enabled_); + if (device_support_resolution_status_.IsOK() && !resolved_acceleration_enabled_ && + acceleration_mode_ == WeightLoadAccelerationMode::Preferred) { LOGS_DEFAULT(WARNING) << "D3D12 accelerated external weights are unavailable; " "using the ordinary WebGPU initializer loading path. " @@ -517,24 +517,24 @@ void D3D12AcceleratedExternalDataLoader::ResolveSupport() const { } common::Status D3D12AcceleratedExternalDataLoader::EnsureFileLoader() const { - if (file_loader_) { + if (file_to_buffer_loader_) { return common::Status::OK(); } return D3D12FileBufferLoader::Create( - context_.WeightLoadingD3D12Device(), file_loader_); + context_.WeightLoadingD3D12Device(), file_to_buffer_loader_); } D3D12AcceleratedExternalDataLoader:: D3D12AcceleratedExternalDataLoader( WebGpuContext& context, std::shared_ptr< - D3D12AcceleratedInitializerState> - state, + D3D12ImportedBufferRegistry> + buffer_registry, WeightLoadAccelerationMode mode) - : context_{context}, state_{std::move(state)}, mode_{mode} { + : context_{context}, buffer_registry_{std::move(buffer_registry)}, acceleration_mode_{mode} { ORT_ENFORCE( - state_ != nullptr, - "D3D12 accelerated allocator state is required."); + buffer_registry_ != nullptr, + "D3D12 imported buffer registry is required."); } D3D12AcceleratedExternalDataLoader:: @@ -547,7 +547,7 @@ bool D3D12AcceleratedExternalDataLoader::CanLoad( ResolveSupport(); return target_memory_info.device == WebGpuDevice && target_memory_info.name == WEBGPU_BUFFER && - enabled_; + resolved_acceleration_enabled_; } bool D3D12AcceleratedExternalDataLoader:: @@ -562,22 +562,22 @@ bool D3D12AcceleratedExternalDataLoader:: const OrtDevice& target_device) const { ResolveSupport(); return target_device == WebGpuDevice && - (enabled_ || - IsWeightLoadAccelerationRequired(mode_)); + (resolved_acceleration_enabled_ || + IsWeightLoadAccelerationRequired(acceleration_mode_)); } common::Status D3D12AcceleratedExternalDataLoader::BeginLoad() const { AbortLoad(); ResolveSupport(); - ORT_RETURN_IF_ERROR(resolved_status_); - if (!enabled_) { + ORT_RETURN_IF_ERROR(device_support_resolution_status_); + if (!resolved_acceleration_enabled_) { return common::Status::OK(); } abort_requested_.store( false, std::memory_order_relaxed); - batch_ = - std::make_unique(); + accelerated_load_batch_ = + std::make_unique(); return common::Status::OK(); } @@ -589,25 +589,25 @@ D3D12AcceleratedExternalDataLoader::PrepareTensor( FileOffsetType data_offset, SafeInt data_length) const { ORT_RETURN_IF_NOT( - batch_ != nullptr && - !batch_->finalized, + accelerated_load_batch_ != nullptr && + !accelerated_load_batch_->finalized, "D3D12 accelerated initializer batch has not been started."); return PrepareTensorForBatch( - *batch_, env, data_file_path, + *accelerated_load_batch_, env, data_file_path, tensor_name, data_offset, data_length); } common::Status D3D12AcceleratedExternalDataLoader::FinalizeLoad( const std::function& is_canceled) const { - if (!enabled_) { + if (!resolved_acceleration_enabled_) { return common::Status::OK(); } ORT_RETURN_IF_NOT( - batch_ != nullptr && - !batch_->finalized, + accelerated_load_batch_ != nullptr && + !accelerated_load_batch_->finalized, "D3D12 accelerated initializer batch has not been started."); - auto& batch = *batch_; + auto& batch = *accelerated_load_batch_; const auto fail_or_fallback = [this, &is_canceled]( const common::Status& status) @@ -620,14 +620,14 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( "D3D12 accelerated initializer loading was canceled."); } if (IsWeightLoadAccelerationRequired( - mode_)) { + acceleration_mode_)) { return status; } LOGS_DEFAULT(WARNING) << "D3D12 accelerated initializer loading failed; " "using the ordinary WebGPU initializer path: " << status.ErrorMessage(); - enabled_ = false; + resolved_acceleration_enabled_ = false; AbortLoad(); return common::Status::OK(); }; @@ -648,7 +648,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( return fail_or_fallback(loader_status); } const auto load_status = LoadBatchToD3D12( - batch, *file_loader_, is_canceled, + batch, *file_to_buffer_loader_, is_canceled, abort_requested_, load_metrics); if (!load_status.IsOK()) { return fail_or_fallback(load_status); @@ -765,9 +765,9 @@ void D3D12AcceleratedExternalDataLoader::AbortLoad() true, std::memory_order_relaxed); ORT_TRY { - if (batch_) { + if (accelerated_load_batch_) { for (auto& tensor : - batch_->tensors) { + accelerated_load_batch_->tensors) { if (!tensor.claimed && tensor.access_started) { ImportedAllocation allocation; @@ -784,7 +784,7 @@ void D3D12AcceleratedExternalDataLoader::AbortLoad() tensor.access_started = false; } } - batch_.reset(); + accelerated_load_batch_.reset(); } } ORT_CATCH(...) { @@ -802,8 +802,8 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( const std::shared_ptr& allocator, Tensor& tensor) const { ORT_RETURN_IF_NOT( - batch_ != nullptr && - batch_->finalized, + accelerated_load_batch_ != nullptr && + accelerated_load_batch_->finalized, "D3D12 accelerated initializer batch has not been finalized."); ORT_RETURN_IF( data_offset < 0, @@ -813,16 +813,16 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( "D3D12 accelerated initializer requires its device allocator."); ORT_RETURN_IF_NOT( allocator.get() == - state_->allocator, + buffer_registry_->allocator, "D3D12 accelerated initializer was passed a different allocator."); const size_t length = static_cast(data_length); const auto file_iterator = - batch_->files.find( + accelerated_load_batch_->files.find( data_file_path.lexically_normal()); ORT_RETURN_IF( - file_iterator == batch_->files.end(), + file_iterator == accelerated_load_batch_->files.end(), "No prepared D3D12 data file matches \"", data_file_path.string(), "\"."); const TensorKey key = MakeKey( @@ -831,16 +831,16 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( static_cast(data_offset), length); const auto iterator = - batch_->tensors_by_key.find(key); + accelerated_load_batch_->tensors_by_key.find(key); ORT_RETURN_IF( iterator == - batch_->tensors_by_key.end(), + accelerated_load_batch_->tensors_by_key.end(), "No prepared D3D12 accelerated initializer matches \"", tensor_name, "\" and the requested file range."); auto& prepared = - batch_->tensors[iterator->second]; + accelerated_load_batch_->tensors[iterator->second]; ORT_RETURN_IF( prepared.claimed, "D3D12 accelerated initializer \"", @@ -881,14 +881,14 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( { std::lock_guard lock{ - state_->mutex}; + buffer_registry_->mutex}; ORT_RETURN_IF( - state_->imported_allocations + buffer_registry_->imported_allocations .find(buffer) != - state_->imported_allocations + buffer_registry_->imported_allocations .end(), "D3D12 accelerated imported buffer was registered twice."); - state_->imported_allocations.emplace( + buffer_registry_->imported_allocations.emplace( buffer, std::move(imported)); } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index b17a9ba9e2f36..83b29ef3f545f 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -33,19 +33,19 @@ common::Status ResolveWeightLoadAccelerationMode( const common::Status& support_status, bool& enabled); -class D3D12AcceleratedInitializerState; -struct D3D12AcceleratedBatch; +class D3D12ImportedBufferRegistry; +struct D3D12AcceleratedLoadBatch; AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( WebGpuContext& context, std::function recording_getter, - std::shared_ptr& out_state); + std::shared_ptr& out_buffer_registry); class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { public: D3D12AcceleratedExternalDataLoader( WebGpuContext& context, - std::shared_ptr state, + std::shared_ptr buffer_registry, WeightLoadAccelerationMode mode); ~D3D12AcceleratedExternalDataLoader() override; @@ -73,14 +73,15 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { common::Status EnsureFileLoader() const; WebGpuContext& context_; - std::shared_ptr state_; - WeightLoadAccelerationMode mode_; - mutable std::once_flag support_once_; - mutable common::Status resolved_status_; - mutable bool enabled_ = false; + std::shared_ptr buffer_registry_; + WeightLoadAccelerationMode acceleration_mode_; + mutable std::once_flag device_support_resolution_once_; + mutable common::Status device_support_resolution_status_; + // Effective decision from mode and device support; cleared on fallback. + mutable bool resolved_acceleration_enabled_ = false; mutable std::atomic abort_requested_{false}; - mutable std::unique_ptr file_loader_; - mutable std::unique_ptr batch_; + mutable std::unique_ptr file_to_buffer_loader_; + mutable std::unique_ptr accelerated_load_batch_; }; } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index caeb5a1a2ca60..d1009143bb18a 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h @@ -42,7 +42,7 @@ class WebGpuContext; class WebGpuProfiler; class GpuBufferAllocator; #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) -class D3D12AcceleratedInitializerState; +class D3D12ImportedBufferRegistry; #endif // Forward declare CapturedCommandInfo which is now defined in webgpu_context.h @@ -160,7 +160,7 @@ class WebGpuExecutionProvider : public IExecutionProvider { #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) webgpu::WeightLoadAccelerationMode weight_load_acceleration_mode_{ webgpu::WeightLoadAccelerationMode::Off}; - std::shared_ptr accelerated_initializer_state_; + std::shared_ptr accelerated_initializer_state_; AllocatorPtr accelerated_initializer_allocator_; #endif std::unordered_map graph_id_to_run_count_; diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index 60216673b300e..d2edc85d965c9 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1441,37 +1441,74 @@ TEST(WebGpuContextTest, D3D12FileLoaderUsesMultipleHeapsAndCommittedResources) { } TEST(WebGpuContextTest, D3D12AcceleratedAllocatorUsesProviderRecording) { + TemporaryDirectory temp_dir{ + ORT_TSTR("webgpu_d3d12_allocator_recording_test")}; + const auto data_path = + std::filesystem::path{temp_dir.Path()} / ORT_TSTR("weights.bin"); + const std::array expected{11, 22, 33, 44}; + { + std::ofstream stream{data_path, std::ios::binary | std::ios::trunc}; + ASSERT_TRUE(stream.good()); + stream.write(reinterpret_cast(expected.data()), + static_cast(sizeof(expected))); + ASSERT_TRUE(stream.good()); + } + auto ep = WebGpuProviderFactoryCreator::Create( WeightLoadAccelerationOptions( - kWeightLoadAcceleration_Preferred)) + kWeightLoadAcceleration_Required)) ->CreateProvider(); ASSERT_NE(ep, nullptr); + auto& context = webgpu::WebGpuContextFactory::GetContext(0); + const auto support_status = + webgpu::CheckD3D12AcceleratedExternalWeightsSupport(context); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + auto loader = ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); auto allocators = ep->CreatePreferredAllocators(); ASSERT_FALSE(allocators.empty()); - auto& allocator = allocators.front(); - void* buffer = allocator->Alloc(16); - ASSERT_NE(buffer, nullptr); + const auto& allocator = allocators.front(); + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->PrepareTensor( + Env::Default(), data_path, "weights", 0, sizeof(expected))); + ASSERT_STATUS_OK(loader->FinalizeLoad([]() { return false; })); auto& webgpu_ep = *static_cast(ep.get()); auto& recording = webgpu_ep.Recording(); - recording.has_unsubmitted_work = true; - allocator->Free(buffer); - constexpr auto usage = - wgpu::BufferUsage::Storage | - wgpu::BufferUsage::CopySrc | - wgpu::BufferUsage::CopyDst | - wgpu::BufferUsage::Indirect; + wgpu::BufferDescriptor destination_description{}; + destination_description.size = sizeof(expected); + destination_description.usage = + wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + auto destination = context.Device().CreateBuffer(&destination_description); + ASSERT_TRUE(destination); + { + Tensor weights{DataTypeImpl::GetType(), TensorShape({4}), + nullptr, allocator}; + ASSERT_STATUS_OK(loader->LoadTensor( + Env::Default(), data_path, "weights", 0, sizeof(expected), + allocator, weights)); + wgpu::Buffer source{ + reinterpret_cast(weights.MutableDataRaw())}; + context.GetCommandEncoder(recording).CopyBufferToBuffer( + source, 0, destination, 0, sizeof(expected)); + } + EXPECT_TRUE(recording.has_unsubmitted_work.load()); + ASSERT_EQ(recording.pending_release_callbacks.size(), 1u); + auto& buffer_manager = webgpu_ep.InitializerBufferManager(); - EXPECT_TRUE(buffer_manager.ExtractCachedBuffers(usage).empty()); - - recording.has_unsubmitted_work = false; - buffer_manager.RefreshPendingBuffers( - recording, webgpu::GraphCaptureState::Default); - auto cached_buffers = - buffer_manager.ExtractCachedBuffers(usage); - EXPECT_EQ(cached_buffers.size(), 1u); - buffer_manager.AbsorbCachedBuffers( - usage, std::move(cached_buffers)); + webgpu::CommandRecordingState unrelated_recording; + ASSERT_STATUS_OK(context.Flush(buffer_manager, unrelated_recording)); + EXPECT_TRUE(recording.has_unsubmitted_work.load()); + EXPECT_EQ(recording.pending_release_callbacks.size(), 1u); + + ASSERT_STATUS_OK(context.Flush(buffer_manager, recording)); + EXPECT_FALSE(recording.has_unsubmitted_work.load()); + EXPECT_TRUE(recording.pending_release_callbacks.empty()); + EXPECT_EQ(ReadBufferWithExternalCommandEncoder<4>(context, destination.Get()), + expected); } TEST(WebGpuContextTest, DuplicateProviderDoesNotRegisterD3D12AcceleratedLoader) { From 5a04a4d269c56e19da33992b8239c9709539ba4e Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Wed, 7 Oct 2026 13:43:20 -0700 Subject: [PATCH 41/46] Refactor D3D12 external data loading methods to improve clarity and consistency in transaction handling --- .../core/framework/external_data_loader.h | 84 ++++++++++--------- .../framework/external_data_loader_manager.cc | 12 +-- .../framework/external_data_loader_manager.h | 4 +- .../core/framework/session_state_utils.cc | 19 +++-- .../core/framework/tensorprotoutils.cc | 9 +- onnxruntime/core/framework/tensorprotoutils.h | 7 +- .../webgpu/d3d12_external_data_loader.cc | 48 +++++------ .../webgpu/d3d12_external_data_loader.h | 16 ++-- .../framework/external_data_loader_test.cc | 48 +++++------ .../providers/webgpu/webgpu_context_test.cc | 44 +++++----- 10 files changed, 148 insertions(+), 143 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index c690a477b23e8..73cee077e1e0b 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -42,59 +42,61 @@ class IExternalDataLoader { // writing into a tensor allocated by the framework. virtual bool CreatesTensorForDevice([[maybe_unused]] const OrtDevice& target_device) const { return false; } - // Begins a batch covering the external initializers for one SessionState - // initialization. ORT describes all loader-created tensors through PrepareTensor() - // before calling FinalizeLoad(), allowing an implementation to coalesce file - // ranges, allocate destination resources together, and perform shared I/O or device - // work before individual tensors become visible through LoadTensor(). Loaders that - // can populate each tensor independently do not need to override the batch hooks. + // Begins one external-weight load transaction for a SessionState initialization. + // A transaction has two phases: + // 1. ORT registers each eligible initializer as it is discovered, then commits the + // complete candidate set to the loader. + // 2. ORT asks the loader for each committed tensor through LoadTensor(). + // EndLoad() ends the transaction after tensor loading or any earlier failure. // - // This is called once before zero or more PrepareTensor() calls and before any - // initializer is loaded through LoadTensor(). Implementations should discard stale - // batch state and initialize any batch-wide resources. Returning an error prevents - // the batch from continuing and causes AbortLoad() to be called. + // Implementations should discard stale transaction state and initialize any + // transaction-wide resources. Returning an error prevents the transaction from + // continuing and causes EndLoad() to be called. virtual common::Status BeginLoad() const { return common::Status::OK(); } - // Adds one external tensor to the current batch. The path, offset, and length + // Registers one external tensor with the active transaction. The path, offset, and length // have already been resolved and validated against the TensorProto by the caller. - // Implementations may record the range, reserve its destination, or start deferred - // work. ORT does not call LoadTensor() during this preparation phase. After - // FinalizeLoad() succeeds, ORT makes a separate initializer-deserialization pass - // and calls LoadTensor() for each tensor still claimed by this loader. Cancellation - // or another initialization error may prevent some or all of those calls. + // This incremental registration lets ORT report candidates as it discovers them + // without retaining a separate core-side candidate collection. The loader owns the + // candidate metadata needed to plan file access, destination allocation, and shared + // device work. ORT does not call LoadTensor() during this phase. // - // PrepareTensor() calls occur after a successful BeginLoad() and before - // FinalizeLoad(). The same file may be referenced by multiple tensors, and a batch - // may contain no tensors. Returning an error abandons the batch and causes - // AbortLoad() to be called. The default implementation performs no preparation. - virtual common::Status PrepareTensor([[maybe_unused]] const Env& env, - [[maybe_unused]] const std::filesystem::path& data_file_path, - [[maybe_unused]] std::string_view tensor_name, - [[maybe_unused]] FileOffsetType data_offset, - [[maybe_unused]] SafeInt data_length) const { + // RegisterLoadCandidate() calls occur after a successful BeginLoad() and before + // CommitLoadCandidates(). The same file may be referenced by multiple tensors, and + // a transaction may contain no candidates. Returning an error ends the transaction. + // The default implementation retains no candidate state. + virtual common::Status RegisterLoadCandidate([[maybe_unused]] const Env& env, + [[maybe_unused]] const std::filesystem::path& data_file_path, + [[maybe_unused]] std::string_view tensor_name, + [[maybe_unused]] FileOffsetType data_offset, + [[maybe_unused]] SafeInt data_length) const { return common::Status::OK(); } - // Completes preparation for the current batch. On success, the loader must - // either make every tensor for which PrepareTensor() returned OK ready for a - // subsequent LoadTensor() call, or make CanLoad() return false for those tensors so - // normal loading can take over. Any asynchronous I/O or device work required to - // make tensors ready must be completed or synchronized here. The implementation - // must keep prepared tensors alive until LoadTensor() transfers their ownership or - // AbortLoad() releases them. This method is called even for an empty batch. + // Commits the complete set of registered candidates before tensor loading begins. + // This is where a loader can coalesce file ranges, allocate destinations, perform + // shared I/O or device work, and synchronize the results. On success, the loader + // must either make each accepted candidate ready for a subsequent LoadTensor() call + // or make CanLoad() return false so normal loading can take over. + // + // ORT then makes a separate initializer-deserialization pass and calls LoadTensor() + // for each tensor still claimed by this loader. The loader must retain unclaimed + // committed tensors until LoadTensor() transfers their ownership or EndLoad() + // releases them. This method is called even when no candidates were registered. // Long-running work should check is_canceled and return MODEL_LOAD_CANCELED when - // cancellation is observed. Returning any error causes AbortLoad() to be called. - // The default implementation completes immediately. - virtual common::Status FinalizeLoad([[maybe_unused]] const std::function& is_canceled) const { + // cancellation is observed. Returning any error causes EndLoad() to be called. + // The default implementation has no work to commit. + virtual common::Status CommitLoadCandidates( + [[maybe_unused]] const std::function& is_canceled) const { return common::Status::OK(); } - // Abandons the active batch. Cancel outstanding work and release resources - // still owned by the batch, without invalidating allocations already transferred by - // successful LoadTensor() calls. This may be called after BeginLoad(), - // PrepareTensor(), FinalizeLoad(), or partial tensor consumption, and may be called - // more than once. Implementations must therefore be noexcept and idempotent. - virtual void AbortLoad() const noexcept {} + // Ends the active transaction. Cancel outstanding work and release candidate + // metadata and resources still owned by the loader, without invalidating allocations + // already transferred by successful LoadTensor() calls. This is called after all + // expected LoadTensor() calls or after any earlier failure or cancellation. It may + // be called more than once, so implementations must be noexcept and idempotent. + virtual void EndLoad() const noexcept {} #endif // Tensor should be allocated with the correct memory info and size. A loader that diff --git a/onnxruntime/core/framework/external_data_loader_manager.cc b/onnxruntime/core/framework/external_data_loader_manager.cc index 86746b4a2126c..688d3d0a90a20 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.cc +++ b/onnxruntime/core/framework/external_data_loader_manager.cc @@ -53,7 +53,7 @@ Status ExternalDataLoaderManager::BeginLoad() const { for (const auto& external_data_loader : external_data_loaders_) { auto status = external_data_loader->BeginLoad(); if (!status.IsOK()) { - AbortLoad(); + EndLoad(); return status; } } @@ -61,11 +61,11 @@ Status ExternalDataLoaderManager::BeginLoad() const { return Status::OK(); } -Status ExternalDataLoaderManager::FinalizeLoad(const std::function& is_canceled) const { +Status ExternalDataLoaderManager::CommitLoadCandidates(const std::function& is_canceled) const { for (const auto& external_data_loader : external_data_loaders_) { - auto status = external_data_loader->FinalizeLoad(is_canceled); + auto status = external_data_loader->CommitLoadCandidates(is_canceled); if (!status.IsOK()) { - AbortLoad(); + EndLoad(); return status; } } @@ -73,9 +73,9 @@ Status ExternalDataLoaderManager::FinalizeLoad(const std::function& is_c return Status::OK(); } -void ExternalDataLoaderManager::AbortLoad() const noexcept { +void ExternalDataLoaderManager::EndLoad() const noexcept { for (const auto& external_data_loader : external_data_loaders_) { - external_data_loader->AbortLoad(); + external_data_loader->EndLoad(); } } #endif diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index 54c8abf5c5db0..259a7f87d3a1b 100644 --- a/onnxruntime/core/framework/external_data_loader_manager.h +++ b/onnxruntime/core/framework/external_data_loader_manager.h @@ -30,8 +30,8 @@ class ExternalDataLoaderManager { const IExternalDataLoader* GetTensorCreator( const OrtDevice& target_device, int32_t tensor_data_type) const; common::Status BeginLoad() const; - common::Status FinalizeLoad(const std::function& is_canceled) const; - void AbortLoad() const noexcept; + common::Status CommitLoadCandidates(const std::function& is_canceled) const; + void EndLoad() const noexcept; #endif private: diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 268dd1dd4286b..fc5d157b0e281 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -409,10 +409,10 @@ common::Status SaveInitializedTensors( #if defined(ENABLE_D3D12_FILE_LOADING) ORT_RETURN_IF_ERROR(external_data_loader_mgr.BeginLoad()); - bool external_data_load_finalized = false; - auto abort_external_data_load = gsl::finally([&]() { - if (!external_data_load_finalized) { - external_data_loader_mgr.AbortLoad(); + bool external_data_load_ended = false; + auto end_external_data_load = gsl::finally([&]() { + if (!external_data_load_ended) { + external_data_loader_mgr.EndLoad(); } }); @@ -432,11 +432,12 @@ common::Status SaveInitializedTensors( exec_plan.GetLocation(entry.first), entry.second->data_type()); if (tensor_creator != nullptr) { ORT_RETURN_IF_ERROR( - utils::PrepareExtDataForTensorFromTensorProto(env, graph_loc, *entry.second, *tensor_creator)); + utils::RegisterExternalDataLoadCandidateFromTensorProto( + env, graph_loc, *entry.second, *tensor_creator)); } } - ORT_RETURN_IF_ERROR(external_data_loader_mgr.FinalizeLoad( + ORT_RETURN_IF_ERROR(external_data_loader_mgr.CommitLoadCandidates( [&session_options]() { return session_options.IsLoadCancellationFlagSet(); })); #endif @@ -528,11 +529,11 @@ common::Status SaveInitializedTensors( } #if defined(ENABLE_D3D12_FILE_LOADING) - // AbortLoad is also the idempotent end-of-batch cleanup hook. All claimed + // EndLoad is the idempotent end-of-batch cleanup hook. All claimed // allocations have transferred ownership, so this only releases metadata // and any unclaimed resources left by a successful batch. - external_data_loader_mgr.AbortLoad(); - external_data_load_finalized = true; + external_data_loader_mgr.EndLoad(); + external_data_load_ended = true; #endif LOGS(logger, INFO) << "Done saving initialized tensors"; return common::Status::OK(); diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 6118622c6325c..cd00b02bfac2f 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -2003,9 +2003,10 @@ Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem: } #if defined(ENABLE_D3D12_FILE_LOADING) -Status PrepareExtDataForTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - const IExternalDataLoader& ext_data_loader) { +Status RegisterExternalDataLoadCandidateFromTensorProto( + const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, + const IExternalDataLoader& ext_data_loader) { ORT_ENFORCE(HasExternalData(tensor_proto)); std::basic_string tensor_proto_dir; @@ -2030,7 +2031,7 @@ Status PrepareExtDataForTensorFromTensorProto(const Env& env, const std::filesys external_data_file_path == onnxruntime::utils::kTensorProtoNativeEndianMemoryAddressTag, "Memory address tag is not supported by custom external data loader."); - return ext_data_loader.PrepareTensor( + return ext_data_loader.RegisterLoadCandidate( env, external_data_file_path, tensor_proto.name(), file_offset, raw_data_safe_len); } diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index b3215c1f2aedb..59bd955c6519e 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -334,9 +334,10 @@ common::Status LoadPrepackedWeightsFromExternalData(const Env& env, const std::f PrepackedWeightsForGraph& prepacked_info); #if defined(ENABLE_D3D12_FILE_LOADING) -common::Status PrepareExtDataForTensorFromTensorProto(const Env& env, const std::filesystem::path& model_path, - const ONNX_NAMESPACE::TensorProto& tensor_proto, - const IExternalDataLoader& ext_data_loader); +common::Status RegisterExternalDataLoadCandidateFromTensorProto( + const Env& env, const std::filesystem::path& model_path, + const ONNX_NAMESPACE::TensorProto& tensor_proto, + const IExternalDataLoader& ext_data_loader); #endif // Convert the AttributeProto from a Constant node into a TensorProto that can be used as an initializer diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 87dcdf8229aae..0819d3b534f0e 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -95,7 +95,7 @@ struct D3D12AcceleratedLoadBatch { std::map files; size_t total_bytes = 0; size_t load_range_count = 0; - bool finalized = false; + bool load_candidates_committed = false; }; namespace { @@ -107,7 +107,7 @@ struct D3D12AcceleratedLoadMetrics { struct CancellationState { const std::function* external = nullptr; - const std::atomic* abort_requested = nullptr; + const std::atomic* load_end_requested = nullptr; }; bool InvokeCancellationNoThrow( @@ -125,8 +125,8 @@ bool InvokeCancellationNoThrow( bool IsCancellationRequested(void* opaque) noexcept { const auto& state = *static_cast(opaque); - if (state.abort_requested != nullptr && - state.abort_requested->load(std::memory_order_relaxed)) { + if (state.load_end_requested != nullptr && + state.load_end_requested->load(std::memory_order_relaxed)) { return true; } if (state.external == nullptr || !*state.external) { @@ -139,9 +139,9 @@ common::Status LoadBatchToD3D12( D3D12AcceleratedLoadBatch& batch, D3D12FileBufferLoader& loader, const std::function& is_canceled, - const std::atomic& abort_requested, + const std::atomic& load_end_requested, D3D12AcceleratedLoadMetrics& metrics) { - if (abort_requested.load(std::memory_order_relaxed) || + if (load_end_requested.load(std::memory_order_relaxed) || InvokeCancellationNoThrow(is_canceled)) { return ORT_MAKE_STATUS( ONNXRUNTIME, MODEL_LOAD_CANCELED, @@ -173,7 +173,7 @@ common::Status LoadBatchToD3D12( auto loaded = std::make_shared(); const CancellationState cancellation_state{ - &is_canceled, &abort_requested}; + &is_canceled, &load_end_requested}; const D3D12FileBufferLoader::CancellationToken cancellation{ IsCancellationRequested, const_cast(&cancellation_state)}; @@ -203,7 +203,7 @@ TensorKey MakeKey(const std::filesystem::path& path, return {path, std::string{name}, offset, length}; } -common::Status PrepareTensorForBatch( +common::Status RegisterLoadCandidateForBatch( D3D12AcceleratedLoadBatch& batch, const Env& env, const std::filesystem::path& data_file_path, @@ -539,7 +539,7 @@ D3D12AcceleratedExternalDataLoader:: D3D12AcceleratedExternalDataLoader:: ~D3D12AcceleratedExternalDataLoader() { - AbortLoad(); + EndLoad(); } bool D3D12AcceleratedExternalDataLoader::CanLoad( @@ -568,13 +568,13 @@ bool D3D12AcceleratedExternalDataLoader:: common::Status D3D12AcceleratedExternalDataLoader::BeginLoad() const { - AbortLoad(); + EndLoad(); ResolveSupport(); ORT_RETURN_IF_ERROR(device_support_resolution_status_); if (!resolved_acceleration_enabled_) { return common::Status::OK(); } - abort_requested_.store( + load_end_requested_.store( false, std::memory_order_relaxed); accelerated_load_batch_ = std::make_unique(); @@ -582,7 +582,7 @@ D3D12AcceleratedExternalDataLoader::BeginLoad() const { } common::Status -D3D12AcceleratedExternalDataLoader::PrepareTensor( +D3D12AcceleratedExternalDataLoader::RegisterLoadCandidate( const Env& env, const std::filesystem::path& data_file_path, std::string_view tensor_name, @@ -590,22 +590,22 @@ D3D12AcceleratedExternalDataLoader::PrepareTensor( SafeInt data_length) const { ORT_RETURN_IF_NOT( accelerated_load_batch_ != nullptr && - !accelerated_load_batch_->finalized, + !accelerated_load_batch_->load_candidates_committed, "D3D12 accelerated initializer batch has not been started."); - return PrepareTensorForBatch( + return RegisterLoadCandidateForBatch( *accelerated_load_batch_, env, data_file_path, tensor_name, data_offset, data_length); } common::Status -D3D12AcceleratedExternalDataLoader::FinalizeLoad( +D3D12AcceleratedExternalDataLoader::CommitLoadCandidates( const std::function& is_canceled) const { if (!resolved_acceleration_enabled_) { return common::Status::OK(); } ORT_RETURN_IF_NOT( accelerated_load_batch_ != nullptr && - !accelerated_load_batch_->finalized, + !accelerated_load_batch_->load_candidates_committed, "D3D12 accelerated initializer batch has not been started."); auto& batch = *accelerated_load_batch_; const auto fail_or_fallback = @@ -628,7 +628,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( "using the ordinary WebGPU initializer path: " << status.ErrorMessage(); resolved_acceleration_enabled_ = false; - AbortLoad(); + EndLoad(); return common::Status::OK(); }; @@ -638,7 +638,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( ONNXRUNTIME, MODEL_LOAD_CANCELED, "D3D12 accelerated initializer loading was canceled.")); } - batch.finalized = true; + batch.load_candidates_committed = true; return common::Status::OK(); } @@ -649,7 +649,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( } const auto load_status = LoadBatchToD3D12( batch, *file_to_buffer_loader_, is_canceled, - abort_requested_, load_metrics); + load_end_requested_, load_metrics); if (!load_status.IsOK()) { return fail_or_fallback(load_status); } @@ -745,7 +745,7 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( return fail_or_fallback(import_status); } const auto import_end = Clock::now(); - batch.finalized = true; + batch.load_candidates_committed = true; LOGS_DEFAULT(VERBOSE) << "WebGPU D3D12 accelerated external initializer load: " @@ -759,9 +759,9 @@ D3D12AcceleratedExternalDataLoader::FinalizeLoad( return common::Status::OK(); } -void D3D12AcceleratedExternalDataLoader::AbortLoad() +void D3D12AcceleratedExternalDataLoader::EndLoad() const noexcept { - abort_requested_.store( + load_end_requested_.store( true, std::memory_order_relaxed); ORT_TRY { @@ -803,8 +803,8 @@ D3D12AcceleratedExternalDataLoader::LoadTensor( Tensor& tensor) const { ORT_RETURN_IF_NOT( accelerated_load_batch_ != nullptr && - accelerated_load_batch_->finalized, - "D3D12 accelerated initializer batch has not been finalized."); + accelerated_load_batch_->load_candidates_committed, + "D3D12 accelerated initializer batch has not been committed."); ORT_RETURN_IF( data_offset < 0, "D3D12 accelerated initializer has a negative file offset."); diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index 83b29ef3f545f..e9fd88a019397 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -53,13 +53,13 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { bool SupportsDataType(int32_t tensor_data_type) const override; bool CreatesTensorForDevice(const OrtDevice& target_device) const override; common::Status BeginLoad() const override; - common::Status PrepareTensor(const Env& env, - const std::filesystem::path& data_file_path, - std::string_view tensor_name, - FileOffsetType data_offset, - SafeInt data_length) const override; - common::Status FinalizeLoad(const std::function& is_canceled) const override; - void AbortLoad() const noexcept override; + common::Status RegisterLoadCandidate(const Env& env, + const std::filesystem::path& data_file_path, + std::string_view tensor_name, + FileOffsetType data_offset, + SafeInt data_length) const override; + common::Status CommitLoadCandidates(const std::function& is_canceled) const override; + void EndLoad() const noexcept override; common::Status LoadTensor(const Env& env, const std::filesystem::path& data_file_path, std::string_view tensor_name, @@ -79,7 +79,7 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { mutable common::Status device_support_resolution_status_; // Effective decision from mode and device support; cleared on fallback. mutable bool resolved_acceleration_enabled_ = false; - mutable std::atomic abort_requested_{false}; + mutable std::atomic load_end_requested_{false}; mutable std::unique_ptr file_to_buffer_loader_; mutable std::unique_ptr accelerated_load_batch_; }; diff --git a/onnxruntime/test/framework/external_data_loader_test.cc b/onnxruntime/test/framework/external_data_loader_test.cc index 51aeb4361e4a2..43b1ff1dbd3b3 100644 --- a/onnxruntime/test/framework/external_data_loader_test.cc +++ b/onnxruntime/test/framework/external_data_loader_test.cc @@ -314,7 +314,7 @@ class BatchLifecycleExternalDataLoader final : public IExternalDataLoader { enum class FailurePoint { None, Begin, - Finalize, + CommitCandidates, }; explicit BatchLifecycleExternalDataLoader(FailurePoint failure_point = FailurePoint::None) @@ -332,20 +332,20 @@ class BatchLifecycleExternalDataLoader final : public IExternalDataLoader { : common::Status::OK(); } - common::Status FinalizeLoad(const std::function&) const override { - ++finalize_count; - return failure_point_ == FailurePoint::Finalize - ? ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "finalize failure") + common::Status CommitLoadCandidates(const std::function&) const override { + ++commit_candidates_count; + return failure_point_ == FailurePoint::CommitCandidates + ? ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "commit failure") : common::Status::OK(); } - void AbortLoad() const noexcept override { - ++abort_count; + void EndLoad() const noexcept override { + ++end_count; } mutable int begin_count = 0; - mutable int finalize_count = 0; - mutable int abort_count = 0; + mutable int commit_candidates_count = 0; + mutable int end_count = 0; private: FailurePoint failure_point_; @@ -362,11 +362,11 @@ TEST(ExternalDataLoaderManagerTest, DefaultBatchLifecycleIsBackwardCompatible) { ExternalDataLoaderManager manager; ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::make_unique())); EXPECT_STATUS_OK(manager.BeginLoad()); - EXPECT_STATUS_OK(manager.FinalizeLoad([]() { return false; })); - manager.AbortLoad(); + EXPECT_STATUS_OK(manager.CommitLoadCandidates([]() { return false; })); + manager.EndLoad(); } -TEST(ExternalDataLoaderManagerTest, BeginFailureAbortsEveryLoader) { +TEST(ExternalDataLoaderManagerTest, BeginFailureEndsEveryLoader) { ExternalDataLoaderManager manager; auto first = std::make_unique(); auto* first_ptr = first.get(); @@ -380,27 +380,27 @@ TEST(ExternalDataLoaderManagerTest, BeginFailureAbortsEveryLoader) { EXPECT_FALSE(manager.BeginLoad().IsOK()); EXPECT_EQ(first_ptr->begin_count, 1); EXPECT_EQ(failing_ptr->begin_count, 1); - EXPECT_EQ(first_ptr->abort_count, 1); - EXPECT_EQ(failing_ptr->abort_count, 1); + EXPECT_EQ(first_ptr->end_count, 1); + EXPECT_EQ(failing_ptr->end_count, 1); } -TEST(ExternalDataLoaderManagerTest, FinalizeFailureAbortsEveryLoader) { +TEST(ExternalDataLoaderManagerTest, CommitCandidatesFailureEndsEveryLoader) { ExternalDataLoaderManager manager; auto failing = std::make_unique( - BatchLifecycleExternalDataLoader::FailurePoint::Finalize); + BatchLifecycleExternalDataLoader::FailurePoint::CommitCandidates); auto* failing_ptr = failing.get(); - auto unfinalized = std::make_unique(); - auto* unfinalized_ptr = unfinalized.get(); + auto uncommitted = std::make_unique(); + auto* uncommitted_ptr = uncommitted.get(); ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(failing))); - ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(unfinalized))); + ASSERT_STATUS_OK(manager.RegisterExternalDataLoader(std::move(uncommitted))); ASSERT_STATUS_OK(manager.BeginLoad()); - EXPECT_FALSE(manager.FinalizeLoad([]() { return false; }).IsOK()); - EXPECT_EQ(failing_ptr->finalize_count, 1); - EXPECT_EQ(unfinalized_ptr->finalize_count, 0); - EXPECT_EQ(failing_ptr->abort_count, 1); - EXPECT_EQ(unfinalized_ptr->abort_count, 1); + EXPECT_FALSE(manager.CommitLoadCandidates([]() { return false; }).IsOK()); + EXPECT_EQ(failing_ptr->commit_candidates_count, 1); + EXPECT_EQ(uncommitted_ptr->commit_candidates_count, 0); + EXPECT_EQ(failing_ptr->end_count, 1); + EXPECT_EQ(uncommitted_ptr->end_count, 1); } #endif diff --git a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc index d2edc85d965c9..7a09b4eead7b5 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -1334,7 +1334,7 @@ TEST(WebGpuContextTest, D3D12AcceleratedCanBeEnabledAfterInitialOffSession) { auto loader = required_ep->GetExternalDataLoader(); ASSERT_NE(loader, nullptr); EXPECT_STATUS_OK(loader->BeginLoad()); - loader->AbortLoad(); + loader->EndLoad(); }); } @@ -1472,9 +1472,9 @@ TEST(WebGpuContextTest, D3D12AcceleratedAllocatorUsesProviderRecording) { ASSERT_FALSE(allocators.empty()); const auto& allocator = allocators.front(); ASSERT_STATUS_OK(loader->BeginLoad()); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", 0, sizeof(expected))); - ASSERT_STATUS_OK(loader->FinalizeLoad([]() { return false; })); + ASSERT_STATUS_OK(loader->CommitLoadCandidates([]() { return false; })); auto& webgpu_ep = *static_cast(ep.get()); auto& recording = webgpu_ep.Recording(); @@ -1597,25 +1597,25 @@ TEST(WebGpuContextTest, D3D12AcceleratedLoadsExternalTensorsAcrossFilesAndRanges } ASSERT_STATUS_OK(loader->BeginLoad()); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", kDataOffset, sizeof(expected))); - const auto duplicate_status = loader->PrepareTensor( + const auto duplicate_status = loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", kDataOffset, sizeof(expected)); EXPECT_FALSE(duplicate_status.IsOK()); EXPECT_NE(duplicate_status.ErrorMessage().find( "Duplicate D3D12 accelerated initializer"), std::string::npos); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "separated", kSeparatedDataOffset, sizeof(separated_expected))); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), other_data_path, "other_weights", kDataOffset, sizeof(other_expected))); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "empty", kSeparatedDataOffset + sizeof(separated_expected), 0)); - ASSERT_STATUS_OK(loader->FinalizeLoad([]() { return false; })); + ASSERT_STATUS_OK(loader->CommitLoadCandidates([]() { return false; })); Tensor weights{DataTypeImpl::GetType(), TensorShape({15}), nullptr, allocator}; @@ -1843,11 +1843,11 @@ TEST(WebGpuContextTest, PreferredD3D12AcceleratedPreservesCancellation) { } ASSERT_STATUS_OK(loader->BeginLoad()); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", 0, sizeof(data))); - const auto status = loader->FinalizeLoad([]() { return true; }); + const auto status = loader->CommitLoadCandidates([]() { return true; }); EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); - loader->AbortLoad(); + loader->EndLoad(); } TEST(WebGpuContextTest, D3D12AcceleratedZeroRequestFinalBatchPreservesCancellation) { @@ -1864,9 +1864,9 @@ TEST(WebGpuContextTest, D3D12AcceleratedZeroRequestFinalBatchPreservesCancellati } ASSERT_STATUS_OK(loader->BeginLoad()); - const auto status = loader->FinalizeLoad([]() { return true; }); + const auto status = loader->CommitLoadCandidates([]() { return true; }); EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); - loader->AbortLoad(); + loader->EndLoad(); } TEST(WebGpuContextTest, D3D12AcceleratedChecksCancellationDuringImport) { @@ -1897,15 +1897,15 @@ TEST(WebGpuContextTest, D3D12AcceleratedChecksCancellationDuringImport) { } ASSERT_STATUS_OK(loader->BeginLoad()); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", 0, sizeof(data))); size_t cancellation_checks = 0; - const auto status = loader->FinalizeLoad([&cancellation_checks]() { + const auto status = loader->CommitLoadCandidates([&cancellation_checks]() { return cancellation_checks++ != 0; }); EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); - loader->AbortLoad(); + loader->EndLoad(); } TEST(WebGpuContextTest, PreferredD3D12AcceleratedFallsBackAfterOperationalFailure) { @@ -1939,10 +1939,10 @@ TEST(WebGpuContextTest, PreferredD3D12AcceleratedFallsBackAfterOperationalFailur } ASSERT_STATUS_OK(loader->BeginLoad()); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", 0, sizeof(data))); ASSERT_TRUE(std::filesystem::remove(data_path)); - EXPECT_STATUS_OK(loader->FinalizeLoad([]() { return false; })); + EXPECT_STATUS_OK(loader->CommitLoadCandidates([]() { return false; })); EXPECT_FALSE(loader->CanLoad(allocator->Info())); } @@ -1974,12 +1974,12 @@ TEST(WebGpuContextTest, PreferredD3D12AcceleratedPreservesConcurrentCancellation } ASSERT_STATUS_OK(loader->BeginLoad()); - ASSERT_STATUS_OK(loader->PrepareTensor( + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( Env::Default(), data_path, "weights", 0, sizeof(data))); ASSERT_TRUE(std::filesystem::remove(data_path)); - const auto status = loader->FinalizeLoad([]() { return true; }); + const auto status = loader->CommitLoadCandidates([]() { return true; }); EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); - loader->AbortLoad(); + loader->EndLoad(); } TEST(WebGpuContextTest, WeightLoadAccelerationModeResolution) { From a9f1a357e9b3865faf01175d930ff362e73ff4c5 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:03:05 -0700 Subject: [PATCH 42/46] Refactor WebGPU external data loader and related components to enhance clarity and remove unused parameters --- .../core/framework/external_data_loader.h | 2 +- .../core/providers/js/external_data_loader.cc | 2 - .../core/providers/js/external_data_loader.h | 2 - .../core/providers/webgpu/allocator.cc | 14 +-- .../webgpu/d3d12_external_data_loader.cc | 4 - .../core/providers/webgpu/data_transfer.cc | 5 -- .../providers/webgpu/external_data_loader.cc | 2 - .../providers/webgpu/external_data_loader.h | 2 - .../core/providers/webgpu/webgpu_context.cc | 5 -- .../core/providers/webgpu/webgpu_context.h | 6 +- .../webgpu/concurrent_context_test.cc | 90 ++----------------- 11 files changed, 14 insertions(+), 120 deletions(-) diff --git a/onnxruntime/core/framework/external_data_loader.h b/onnxruntime/core/framework/external_data_loader.h index 73cee077e1e0b..1e8ca0dd8f68f 100644 --- a/onnxruntime/core/framework/external_data_loader.h +++ b/onnxruntime/core/framework/external_data_loader.h @@ -102,7 +102,7 @@ class IExternalDataLoader { // Tensor should be allocated with the correct memory info and size. A loader that // creates tensors for the target device replaces it with one backed by memory // owned through allocator. -#if defined(ENABLE_D3D12_FILE_LOADING) || defined(__wasm__) +#if defined(ENABLE_D3D12_FILE_LOADING) virtual common::Status LoadTensor([[maybe_unused]] const Env& env, [[maybe_unused]] const std::filesystem::path& data_file_path, [[maybe_unused]] std::string_view tensor_name, diff --git a/onnxruntime/core/providers/js/external_data_loader.cc b/onnxruntime/core/providers/js/external_data_loader.cc index f5336d1754bad..193b373cf3696 100644 --- a/onnxruntime/core/providers/js/external_data_loader.cc +++ b/onnxruntime/core/providers/js/external_data_loader.cc @@ -20,10 +20,8 @@ bool ExternalDataLoader::CanLoad(const OrtMemoryInfo& target_memory_info) const common::Status ExternalDataLoader::LoadTensor(const Env& env, const std::filesystem::path& data_file_path, - std::string_view /*tensor_name*/, FileOffsetType data_offset, SafeInt data_length, - const std::shared_ptr& /*allocator*/, Tensor& tensor) const { ExternalDataLoadType load_type; if (tensor.Location().device.Type() == OrtDevice::CPU) { diff --git a/onnxruntime/core/providers/js/external_data_loader.h b/onnxruntime/core/providers/js/external_data_loader.h index 89fddb3d056ac..5f35ed62bbcc1 100644 --- a/onnxruntime/core/providers/js/external_data_loader.h +++ b/onnxruntime/core/providers/js/external_data_loader.h @@ -17,10 +17,8 @@ class ExternalDataLoader : public IExternalDataLoader { common::Status LoadTensor(const Env& env, const std::filesystem::path& data_file_path, - std::string_view tensor_name, FileOffsetType data_offset, SafeInt data_length, - const std::shared_ptr& allocator, Tensor& tensor) const override; }; diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index dc973c801f7aa..46d81402cf8b7 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include #include #include @@ -50,24 +49,19 @@ void* GpuBufferAllocator::Allocate(size_t size, CommandRecordingState& recording return nullptr; } - std::lock_guard lock{recording.mutex}; - const auto& buffer_manager = buffer_manager_getter_(); wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; - auto buffer = buffer_manager.Create(recording, size, usage, initialize_to_zero_, submit_zero_initialize); + auto buffer = buffer_manager_getter_().Create(recording, size, usage, initialize_to_zero_, + submit_zero_initialize); num_allocs_.fetch_add(1, std::memory_order_relaxed); return buffer; } void GpuBufferAllocator::Free(void* p) { if (p != nullptr) { - auto* recording = recording_getter_ ? &recording_getter_() : nullptr; - std::unique_lock lock; - if (recording != nullptr) { - lock = std::unique_lock{recording->mutex}; - } - buffer_manager_getter_().Release(static_cast(p), recording); + buffer_manager_getter_().Release(static_cast(p), + recording_getter_ ? &recording_getter_() : nullptr); num_allocs_.fetch_sub(1, std::memory_order_relaxed); } } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 0819d3b534f0e..27f33a68a446a 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -408,8 +408,6 @@ class D3D12AcceleratedWebGpuAllocator final wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; auto& recording = recording_getter_(); - std::lock_guard lock{ - recording.mutex}; return context_.InitializerBufferManager().Create( recording, size, usage); } @@ -420,8 +418,6 @@ class D3D12AcceleratedWebGpuAllocator final } auto& recording = recording_getter_(); - std::lock_guard - recording_lock{recording.mutex}; std::unique_ptr imported; { std::lock_guard lock{ diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 2ad2ac9621b28..4411c922a6fc3 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -3,8 +3,6 @@ #include "core/providers/webgpu/data_transfer.h" -#include - #include "core/providers/webgpu/buffer_manager.h" #include "core/providers/webgpu/webgpu_context.h" @@ -17,9 +15,6 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data, bool dst_is_gpu, size_t bytes) const { auto& command_state = recording_; - // Public allocator calls can share this recorder; serialize copy encoding, - // not just submission, against allocation and release. - std::lock_guard lock{command_state.mutex}; if (bytes > 0) { if (dst_is_gpu) { if (src_is_gpu) { diff --git a/onnxruntime/core/providers/webgpu/external_data_loader.cc b/onnxruntime/core/providers/webgpu/external_data_loader.cc index 51bc230805b54..6da9598b146f5 100644 --- a/onnxruntime/core/providers/webgpu/external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/external_data_loader.cc @@ -18,10 +18,8 @@ bool ExternalDataLoader::CanLoad(const OrtMemoryInfo& target_memory_info) const common::Status ExternalDataLoader::LoadTensor(const Env& env, const std::filesystem::path& data_file_path, - std::string_view /*tensor_name*/, FileOffsetType data_offset, SafeInt data_length, - const std::shared_ptr& /*allocator*/, Tensor& tensor) const { ExternalDataLoadType load_type; if (tensor.Location().device.Type() == OrtDevice::CPU) { diff --git a/onnxruntime/core/providers/webgpu/external_data_loader.h b/onnxruntime/core/providers/webgpu/external_data_loader.h index 15a25bca3dd32..7ced4e930bf7a 100644 --- a/onnxruntime/core/providers/webgpu/external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/external_data_loader.h @@ -19,10 +19,8 @@ class ExternalDataLoader : public IExternalDataLoader { common::Status LoadTensor(const Env& env, const std::filesystem::path& data_file_path, - std::string_view tensor_name, FileOffsetType data_offset, SafeInt data_length, - const std::shared_ptr& allocator, Tensor& tensor) const override; }; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 0643654634da5..ed7b097cde45e 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -632,7 +632,6 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra const webgpu::BufferManager& buffer_mgr = ComputeContextBase::BufferManagerAccessor::Get(context); CommandRecordingState& recording = ComputeContextBase::BufferManagerAccessor::GetRecording(context); - std::lock_guard lock{recording.mutex}; // validate inputs and outputs are on WebGPU buffers if (ValidationMode() >= ValidationMode::Basic) { @@ -1391,7 +1390,6 @@ Status WebGpuContext::FlushAndWaitChecked(const webgpu::BufferManager& buffer_mg Status WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr, CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; // Graph runs and prepacking can release allocator/uniform buffers into the shared manager. // Retire this recording in both managers, keeping graph cache policy local to the active one. const auto refresh_pending_buffers = [&]() { @@ -1539,7 +1537,6 @@ void WebGpuContext::DispatchCommand(const webgpu::CapturedCommandInfo& command, void WebGpuContext::CaptureBegin(std::vector* captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "CaptureBegin with external storage"; // Flush any pending commands before we change the status ORT_THROW_IF_ERROR(Flush(buffer_manager, recording)); @@ -1556,7 +1553,6 @@ void WebGpuContext::CaptureBegin(std::vector* captu void WebGpuContext::Replay(const std::vector& captured_commands, const webgpu::BufferManager& buffer_manager, CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "Replay with external storage"; recording.graph_capture_state = GraphCaptureState::Replaying; // Replay all captured commands from the provided vector @@ -1587,7 +1583,6 @@ void WebGpuContext::Replay(const std::vector& captu } void WebGpuContext::CaptureEnd(CommandRecordingState& recording) { - std::lock_guard lock{recording.mutex}; LOGS_DEFAULT(VERBOSE) << "CaptureEnd"; recording.graph_capture_state = GraphCaptureState::Default; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 423d5229f11f2..e02777b1d78b3 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -116,13 +116,13 @@ struct CapturedCommandInfo { std::optional pending_kernel_info; }; -// Serializes callers that share a command recording timeline, including Run, public -// allocator use, data transfer, and graph capture operations. +// A command recording timeline has one caller at a time. ORT serializes Run on a Session; +// plugin streamless operations use independent encoders. BufferManager synchronizes buffer +// releases from concurrent allocator callers separately from command recording. struct CommandRecordingState { CommandRecordingState() = default; ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CommandRecordingState); - std::recursive_mutex mutex; wgpu::CommandEncoder command_encoder; wgpu::ComputePassEncoder compute_pass_encoder; uint32_t num_pending_dispatches = 0; diff --git a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc index 0118613df8c29..47f57bec1d4db 100644 --- a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc +++ b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc @@ -450,85 +450,7 @@ TEST_F(WebGpuConcurrentContextTest, ColdAndWarmSessionsRunConcurrently) { << throughput_efficiency << std::endl; } -// Case F: the public session allocator wraps this same internal allocator. Allocations made -// concurrently with Run must not race the session's command recording state. -TEST_F(WebGpuConcurrentContextTest, SessionAllocatorAndRunConcurrently) { - constexpr int kIters = 40; - auto session = MakeSession(); - auto allocator = session->GetAllocator(OrtMemoryInfo(WEBGPU_BUFFER, - OrtAllocatorType::OrtDeviceAllocator, - webgpu::WebGpuDevice, - OrtMemTypeDefault)); - ASSERT_NE(allocator, nullptr); - - ErrorSink sink; - std::barrier start{2}; - std::thread runner([&]() { - start.arrive_and_wait(); - RunLoop(*session, kIters, sink, "F.run"); - }); - std::thread allocator_user([&]() { - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - Tensor tensor(DataTypeImpl::GetType(), TensorShape{kNumElements}, allocator); - ASSERT_NE(tensor.MutableDataRaw(), nullptr); - } - } catch (const std::exception& e) { - sink.Record(std::string("F.allocator threw: ") + e.what()); - } - }); - runner.join(); - allocator_user.join(); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case G: exercise the plugin Env allocator's shared implementation and getter ownership in a -// native build. Its private command state must remain thread-safe across callers. -TEST_F(WebGpuConcurrentContextTest, SharedAllocatorMultiThreadCreateTensor) { - constexpr int kThreads = 4; - constexpr int kIters = 60; - auto& context = webgpu::WebGpuContextFactory::GetContext(0); - auto context_ref = std::shared_ptr(&context, [](webgpu::WebGpuContext*) {}); - auto recording = std::make_shared(); - std::weak_ptr recording_lifetime = recording; - auto allocator = std::make_shared( - [context_ref = std::move(context_ref)]() -> const webgpu::BufferManager& { - return context_ref->BufferManager(); - }, - [recording = std::move(recording)]() -> webgpu::CommandRecordingState& { return *recording; }, - false, - []() { return true; }); - - ErrorSink sink; - std::barrier start{kThreads}; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - Tensor tensor(DataTypeImpl::GetType(), TensorShape{1024 + t * 16}, allocator); - ASSERT_NE(tensor.MutableDataRaw(), nullptr); - } - } catch (const std::exception& e) { - sink.Record("G.thread" + std::to_string(t) + " threw: " + e.what()); - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); - EXPECT_FALSE(recording_lifetime.expired()); - AllocatorStats stats{}; - allocator->GetStats(&stats); - EXPECT_EQ(stats.num_allocs, 0); - allocator.reset(); - EXPECT_TRUE(recording_lifetime.expired()); -} - -// Case H: OrtEnv owns one data-transfer implementation per EP factory. Concurrent CopyTensors +// Case F: OrtEnv owns one data-transfer implementation per EP factory. Concurrent CopyTensors // calls must not encode and flush through the same recording state simultaneously. TEST_F(WebGpuConcurrentContextTest, SharedDataTransferMultiThreadCopy) { constexpr int kThreads = 4; @@ -561,11 +483,11 @@ TEST_F(WebGpuConcurrentContextTest, SharedDataTransferMultiThreadCopy) { ORT_THROW_IF_ERROR(data_transfer.CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, output.size() * sizeof(float))); if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { - sink.Record("H.thread" + std::to_string(t) + " copied incorrect data"); + sink.Record("F.thread" + std::to_string(t) + " copied incorrect data"); } } } catch (const std::exception& e) { - sink.Record("H.thread" + std::to_string(t) + " threw: " + e.what()); + sink.Record("F.thread" + std::to_string(t) + " threw: " + e.what()); } }); } @@ -574,7 +496,7 @@ TEST_F(WebGpuConcurrentContextTest, SharedDataTransferMultiThreadCopy) { ASSERT_FALSE(sink.Failed()) << sink.FirstError(); } -// Case I: separate data-transfer objects share the context-level BufferManager but own distinct +// Case G: separate data-transfer objects share the context-level BufferManager but own distinct // command recording timelines. Their command encoders and pending buffers must remain isolated. TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { constexpr int kThreads = 4; @@ -608,11 +530,11 @@ TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { ORT_THROW_IF_ERROR(data_transfers[t]->CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, output.size() * sizeof(float))); if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { - sink.Record("I.thread" + std::to_string(t) + " copied incorrect data"); + sink.Record("G.thread" + std::to_string(t) + " copied incorrect data"); } } } catch (const std::exception& e) { - sink.Record("I.thread" + std::to_string(t) + " threw: " + e.what()); + sink.Record("G.thread" + std::to_string(t) + " threw: " + e.what()); } }); } From 1ae02b08844668ed99b061f7dfd1964e07b33175 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:11:32 -0700 Subject: [PATCH 43/46] Fix WebGPU concurrent data transfer test Remove the test that shares one unsupported low-level command recording state across threads. Retain the independent data transfer coverage that matches the supported ownership model. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../webgpu/concurrent_context_test.cc | 52 ++----------------- 1 file changed, 3 insertions(+), 49 deletions(-) diff --git a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc index 47f57bec1d4db..7109b957a873b 100644 --- a/onnxruntime/test/providers/webgpu/concurrent_context_test.cc +++ b/onnxruntime/test/providers/webgpu/concurrent_context_test.cc @@ -450,53 +450,7 @@ TEST_F(WebGpuConcurrentContextTest, ColdAndWarmSessionsRunConcurrently) { << throughput_efficiency << std::endl; } -// Case F: OrtEnv owns one data-transfer implementation per EP factory. Concurrent CopyTensors -// calls must not encode and flush through the same recording state simultaneously. -TEST_F(WebGpuConcurrentContextTest, SharedDataTransferMultiThreadCopy) { - constexpr int kThreads = 4; - constexpr int kIters = 30; - constexpr size_t kElements = 4096; - auto& context = webgpu::WebGpuContextFactory::GetContext(0); - webgpu::CommandRecordingState recording; - webgpu::DataTransferImpl data_transfer(context.BufferManager(), recording); - - std::array gpu_buffers; - for (auto& buffer : gpu_buffers) { - wgpu::BufferDescriptor desc{}; - desc.size = kElements * sizeof(float); - desc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; - buffer = context.Device().CreateBuffer(&desc); - } - - ErrorSink sink; - std::barrier start{kThreads}; - std::vector threads; - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&, t]() { - std::vector input(kElements, static_cast(t + 1)); - std::vector output(kElements); - start.arrive_and_wait(); - try { - for (int i = 0; i < kIters && !sink.Failed(); ++i) { - ORT_THROW_IF_ERROR(data_transfer.CopyTensor(input.data(), false, gpu_buffers[t].Get(), true, - input.size() * sizeof(float))); - ORT_THROW_IF_ERROR(data_transfer.CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, - output.size() * sizeof(float))); - if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { - sink.Record("F.thread" + std::to_string(t) + " copied incorrect data"); - } - } - } catch (const std::exception& e) { - sink.Record("F.thread" + std::to_string(t) + " threw: " + e.what()); - } - }); - } - JoinAll(threads); - - ASSERT_FALSE(sink.Failed()) << sink.FirstError(); -} - -// Case G: separate data-transfer objects share the context-level BufferManager but own distinct +// Case F: separate data-transfer objects share the context-level BufferManager but own distinct // command recording timelines. Their command encoders and pending buffers must remain isolated. TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { constexpr int kThreads = 4; @@ -530,11 +484,11 @@ TEST_F(WebGpuConcurrentContextTest, IndependentDataTransfersMultiThreadCopy) { ORT_THROW_IF_ERROR(data_transfers[t]->CopyTensor(gpu_buffers[t].Get(), true, output.data(), false, output.size() * sizeof(float))); if (!std::all_of(output.begin(), output.end(), [&](float value) { return value == input[0]; })) { - sink.Record("G.thread" + std::to_string(t) + " copied incorrect data"); + sink.Record("F.thread" + std::to_string(t) + " copied incorrect data"); } } } catch (const std::exception& e) { - sink.Record("G.thread" + std::to_string(t) + " threw: " + e.what()); + sink.Record("F.thread" + std::to_string(t) + " threw: " + e.what()); } }); } From 32c852a3a1d62b66cf284ba425555a1e3c23c1b1 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:42:02 -0700 Subject: [PATCH 44/46] Adapt D3D12 loading to WebGPU device IDs Propagate the WebGPU context ID through the accelerated allocator and external loader so their memory identity matches the current main branch interfaces. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../webgpu/d3d12_external_data_loader.cc | 16 +++++++++++----- .../webgpu/d3d12_external_data_loader.h | 3 +++ .../webgpu/webgpu_execution_provider.cc | 3 ++- .../providers/webgpu/webgpu_provider_factory.cc | 2 -- 4 files changed, 16 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc index 27f33a68a446a..08d197a485139 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -381,6 +381,7 @@ class D3D12AcceleratedWebGpuAllocator final public: D3D12AcceleratedWebGpuAllocator( WebGpuContext& context, + int context_id, std::function recording_getter, std::shared_ptr @@ -389,7 +390,7 @@ class D3D12AcceleratedWebGpuAllocator final OrtMemoryInfo( WEBGPU_BUFFER, OrtAllocatorType::OrtReadOnlyAllocator, - WebGpuDevice, + WebGpuDevice(context_id), OrtMemTypeDefault)), context_{context}, recording_getter_{ @@ -474,6 +475,7 @@ class D3D12AcceleratedWebGpuAllocator final AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( WebGpuContext& context, + int context_id, std::function recording_getter, std::shared_ptr& @@ -483,7 +485,7 @@ AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( auto allocator = std::make_shared< D3D12AcceleratedWebGpuAllocator>( - context, std::move(recording_getter), buffer_registry); + context, context_id, std::move(recording_getter), buffer_registry); buffer_registry->allocator = allocator.get(); out_buffer_registry = buffer_registry; return allocator; @@ -523,11 +525,15 @@ common::Status D3D12AcceleratedExternalDataLoader::EnsureFileLoader() const { D3D12AcceleratedExternalDataLoader:: D3D12AcceleratedExternalDataLoader( WebGpuContext& context, + int context_id, std::shared_ptr< D3D12ImportedBufferRegistry> buffer_registry, WeightLoadAccelerationMode mode) - : context_{context}, buffer_registry_{std::move(buffer_registry)}, acceleration_mode_{mode} { + : context_{context}, + context_id_{context_id}, + buffer_registry_{std::move(buffer_registry)}, + acceleration_mode_{mode} { ORT_ENFORCE( buffer_registry_ != nullptr, "D3D12 imported buffer registry is required."); @@ -541,7 +547,7 @@ D3D12AcceleratedExternalDataLoader:: bool D3D12AcceleratedExternalDataLoader::CanLoad( const OrtMemoryInfo& target_memory_info) const { ResolveSupport(); - return target_memory_info.device == WebGpuDevice && + return target_memory_info.device == WebGpuDevice(context_id_) && target_memory_info.name == WEBGPU_BUFFER && resolved_acceleration_enabled_; } @@ -557,7 +563,7 @@ bool D3D12AcceleratedExternalDataLoader:: CreatesTensorForDevice( const OrtDevice& target_device) const { ResolveSupport(); - return target_device == WebGpuDevice && + return target_device == WebGpuDevice(context_id_) && (resolved_acceleration_enabled_ || IsWeightLoadAccelerationRequired(acceleration_mode_)); } diff --git a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h index e9fd88a019397..13a5969ff4866 100644 --- a/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -38,6 +38,7 @@ struct D3D12AcceleratedLoadBatch; AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( WebGpuContext& context, + int context_id, std::function recording_getter, std::shared_ptr& out_buffer_registry); @@ -45,6 +46,7 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { public: D3D12AcceleratedExternalDataLoader( WebGpuContext& context, + int context_id, std::shared_ptr buffer_registry, WeightLoadAccelerationMode mode); ~D3D12AcceleratedExternalDataLoader() override; @@ -73,6 +75,7 @@ class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { common::Status EnsureFileLoader() const; WebGpuContext& context_; + int context_id_; std::shared_ptr buffer_registry_; WeightLoadAccelerationMode acceleration_mode_; mutable std::once_flag device_support_resolution_once_; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 271c8b1db64f2..dbf90a0aee31c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -636,6 +636,7 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, accelerated_initializer_allocator_ = CreateD3D12AcceleratedWebGpuAllocator( context_, + context_id_, [this]() -> webgpu::CommandRecordingState& { return Recording(); }, accelerated_initializer_state_); } @@ -813,7 +814,7 @@ std::unique_ptr WebGpuExecutionProvider::GetEx } return std::make_unique( - context_, accelerated_initializer_state_, + context_, context_id_, accelerated_initializer_state_, weight_load_acceleration_mode_); } #endif diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 2d7434269cd14..dffe54f41a3d8 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -251,8 +251,6 @@ WebGpuContextConfig ParseWebGpuContextConfig( config.context_id >= 0 && config.context_id <= std::numeric_limits::max(), "Invalid deviceId value: ", context_id_str, ". Must be an integer in the range 0 to ", std::numeric_limits::max(), "."); - ORT_ENFORCE(config.context_id != kDeviceFreeDefaultContextId, - "WebGPU device ID ", kDeviceFreeDefaultContextId, " is reserved for internal use."); } if (std::string adapter_index_str; From da8d322f338aaf6818afbff79d1c5922b7aafbe7 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:51:02 -0700 Subject: [PATCH 45/46] Restore WebGPU device error propagation Preserve uncaptured error and device-loss state for checked completion, and apply the required C++ formatting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/webgpu/webgpu_context.cc | 25 +++++++++++++++---- .../webgpu/webgpu_execution_provider.cc | 12 ++++----- .../webgpu/webgpu_provider_factory.cc | 2 +- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index ed7b097cde45e..2a7c4da531e4f 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -308,20 +308,35 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { wgpu::Limits required_limits = GetRequiredLimits(adapter); device_desc.requiredLimits = &required_limits; - // TODO: revise temporary error handling device_desc.SetUncapturedErrorCallback( // Note: Don't throw from a Dawn callback. [](const wgpu::Device& /*device*/, wgpu::ErrorType type, - wgpu::StringView message) noexcept { + wgpu::StringView message, DeviceErrorState* state) noexcept { + { + std::lock_guard lock{state->mutex}; + if (state->status.IsOK()) { + state->status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU device error(", int(type), + "): ", std::string_view{message}); + } + } if (logging::LoggingManager::HasDefaultLogger()) { LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; } - }); - // TODO: revise temporary device lost handling + }, + device_error_state_.get()); + // Dawn stops uncaptured callbacks before invoking the lost callback, which retains their state. device_desc.SetDeviceLostCallback( wgpu::CallbackMode::AllowSpontaneous, // Note: Don't throw from a Dawn callback. - [](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, wgpu::StringView message) noexcept { + [state = device_error_state_](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, + wgpu::StringView message) noexcept { + { + std::lock_guard lock{state->mutex}; + if (state->status.IsOK()) { + state->status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU device lost (", int(reason), + "): ", std::string_view{message}); + } + } if (logging::LoggingManager::HasDefaultLogger()) { LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; } diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index dbf90a0aee31c..ec2ed34439ab6 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -673,18 +673,18 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, std::vector WebGpuExecutionProvider::CreatePreferredAllocators() { const bool device_free = !context_.HasDevice(); return { - // allocator for initializers + // allocator for initializers #if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) context_.HasDevice() && accelerated_initializer_allocator_ != nullptr ? accelerated_initializer_allocator_ : #endif - CreateWebGpuAllocator( - context_id_, - device_free, - [this]() -> const webgpu::BufferManager& { return InitializerBufferManager(); }, - [this]() -> webgpu::CommandRecordingState& { return Recording(); }, true), + CreateWebGpuAllocator( + context_id_, + device_free, + [this]() -> const webgpu::BufferManager& { return InitializerBufferManager(); }, + [this]() -> webgpu::CommandRecordingState& { return Recording(); }, true), // default allocator CreateWebGpuAllocator( context_id_, diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index dffe54f41a3d8..843bb5cd8893b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -246,7 +246,7 @@ WebGpuContextConfig ParseWebGpuContextConfig( if (std::string context_id_str; config_options.TryGetConfigEntry(kDeviceId, context_id_str)) { const auto result = std::from_chars(context_id_str.data(), context_id_str.data() + context_id_str.size(), - config.context_id); + config.context_id); ORT_ENFORCE(result.ec == std::errc{} && result.ptr == context_id_str.data() + context_id_str.size() && config.context_id >= 0 && config.context_id <= std::numeric_limits::max(), "Invalid deviceId value: ", context_id_str, ". Must be an integer in the range 0 to ", From 4b92e7db7718263c0ea23a5a41a6391cc270d967 Mon Sep 17 00:00:00 2001 From: Sushanth Rajasankar <44513542+sushraja-msft@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:50:06 -0700 Subject: [PATCH 46/46] Preserve in-memory external data loaders Route memory-backed external data through ordinary loaders while excluding tensor-creating file loaders such as D3D12 acceleration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/framework/session_state_utils.cc | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index fc5d157b0e281..3dd06758ca9ef 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -99,10 +99,18 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st const auto& memory_info = (alloc != nullptr) ? alloc->Info() : memory_buffer->GetAllocInfo(); const auto device = memory_info.device; - if (utils::HasExternalData(tensor_proto) && - !utils::HasExternalDataInMemory(tensor_proto)) { + if (utils::HasExternalData(tensor_proto)) { auto external_data_loader = external_data_loader_mgr.GetExternalDataLoader(memory_info, tensor_proto.data_type()); +#if defined(ENABLE_D3D12_FILE_LOADING) + // Tensor-creating file loaders cannot consume memory-backed external data. + // Other loaders, including WebAssembly loaders, must still receive it. + if (utils::HasExternalDataInMemory(tensor_proto) && + external_data_loader != nullptr && + external_data_loader->CreatesTensorForDevice(device)) { + external_data_loader = nullptr; + } +#endif if (external_data_loader) { #if defined(ENABLE_D3D12_FILE_LOADING) if (external_data_loader->CreatesTensorForDevice(device)) {