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..c80d73aeacb90 --- /dev/null +++ b/cmake/onnxruntime_d3d12_file_loader.cmake @@ -0,0 +1,35 @@ +# 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") + +if(NOT onnxruntime_BUILD_SHARED_LIB) + install( + TARGETS onnxruntime_d3d12_file_loader + EXPORT ${PROJECT_NAME}Targets + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endif() 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..f2b33c91fab90 100644 --- a/js/node/test/api/simple-api-tests.ts +++ b/js/node/test/api/simple-api-tests.ts @@ -88,6 +88,37 @@ 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'); + let unavailableProviderError: string | undefined; + + for (const weightLoadAcceleration of values) { + try { + const session = await InferenceSession.create(model, { + executionProviders: [{ name: 'webgpu', weightLoadAcceleration }], + }); + await session.release(); + } catch (error) { + 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') { + if (message !== unavailableProviderError) { + assert.match(message, /weightLoadAcceleration|disk-to-GPU|DXGI|D3D12/); + } + } else { + assert.doesNotMatch(message, /weightLoadAcceleration|disk-to-GPU|DXGI|D3D12/); + } + } + } + }); + 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..1e8ca0dd8f68f 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,92 @@ 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 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. + // + // 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(); } + + // 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. + // 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. + // + // 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(); + } + + // 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 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(); + } + + // 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 + // creates tensors for the target device replaces it with one backed by memory + // owned through allocator. +#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, + [[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..688d3d0a90a20 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,59 @@ 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; +} + +#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_) { + 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()) { + EndLoad(); + return status; + } + } + + return Status::OK(); +} + +Status ExternalDataLoaderManager::CommitLoadCandidates(const std::function& is_canceled) const { + for (const auto& external_data_loader : external_data_loaders_) { + auto status = external_data_loader->CommitLoadCandidates(is_canceled); + if (!status.IsOK()) { + EndLoad(); + return status; + } + } + + return Status::OK(); +} + +void ExternalDataLoaderManager::EndLoad() const noexcept { + for (const auto& external_data_loader : external_data_loaders_) { + external_data_loader->EndLoad(); + } +} +#endif + } // namespace onnxruntime diff --git a/onnxruntime/core/framework/external_data_loader_manager.h b/onnxruntime/core/framework/external_data_loader_manager.h index c2bcd1c9034e7..259a7f87d3a1b 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" @@ -22,6 +24,16 @@ 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; +#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 CommitLoadCandidates(const std::function& is_canceled) const; + void EndLoad() 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 db988873539ad..3dd06758ca9ef 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -100,12 +100,34 @@ 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 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 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 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 +#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)); + } 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 +352,18 @@ 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); +#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)); if (trace_allocation) { // can not trace string tensor, and they exist only on CPU @@ -350,6 +382,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)); } @@ -375,6 +415,40 @@ 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_ended = false; + auto end_external_data_load = gsl::finally([&]() { + if (!external_data_load_ended) { + external_data_loader_mgr.EndLoad(); + } + }); + + 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::RegisterExternalDataLoadCandidateFromTensorProto( + env, graph_loc, *entry.second, *tensor_creator)); + } + } + + ORT_RETURN_IF_ERROR(external_data_loader_mgr.CommitLoadCandidates( + [&session_options]() { return session_options.IsLoadCancellationFlagSet(); })); +#endif + // 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 +536,13 @@ common::Status SaveInitializedTensors( #endif } +#if defined(ENABLE_D3D12_FILE_LOADING) + // 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.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 52a4cfe90a10b..cd00b02bfac2f 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -1958,7 +1958,12 @@ 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) { +#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. @@ -1980,6 +1985,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 +2002,41 @@ Status LoadExtDataToTensorFromTensorProto(const Env& env, const std::filesystem: #endif } +#if defined(ENABLE_D3D12_FILE_LOADING) +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; + 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.RegisterLoadCandidate( + 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..59bd955c6519e 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,13 @@ 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 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 // 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..bc3aa37ceeb16 --- /dev/null +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.cc @@ -0,0 +1,1159 @@ +// 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 +#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; +// 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( + ONNXRUNTIME, FAIL, operation, " failed with HRESULT 0x", + std::setfill('0'), std::uppercase, std::hex, std::setw(8), + 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 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, + "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::FileReadRegion { + uint64_t begin = 0; + uint64_t end = 0; +}; + +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 buffer_source_indices; + std::vector read_regions; +}; + +struct D3D12FileBufferLoader::UploadStagingSlot { + ComPtr resource; + ComPtr allocator; + ComPtr command_list; + wil::unique_handle file_read_complete_event; + void* mapped = nullptr; + OVERLAPPED overlapped{}; + HANDLE read_file = INVALID_HANDLE_VALUE; + uint64_t file_offset = 0; + DWORD requested_read_bytes = 0; + uint64_t fence_value = 0; + bool read_active = false; +}; + +struct D3D12FileBufferLoader::AllocationResult { + common::Status status; + BufferCollection batch; +}; + +D3D12FileBufferLoader::D3D12FileBufferLoader( + ID3D12Device* device, const Config& config) noexcept + : device_(device), config_(config) { +} + +D3D12FileBufferLoader::~D3D12FileBufferLoader() { + 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& 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& staging_slot : upload_staging_slots_) { + if (staging_slot.resource && staging_slot.mapped != nullptr) { + staging_slot.resource->Unmap(0, nullptr); + staging_slot.mapped = nullptr; + } + } +} + +common::Status D3D12FileBufferLoader::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); + } + + copy_fence_complete_event_.reset(CreateEventExW( + nullptr, nullptr, 0, EVENT_MODIFY_STATE | SYNCHRONIZE)); + if (!copy_fence_complete_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); + 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(&staging_slot.resource)); + if (FAILED(hr)) { + return HResultError( + "ID3D12Device::CreateCommittedResource(upload)", hr); + } + + 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(&staging_slot.allocator)); + if (FAILED(hr)) { + return HResultError( + "ID3D12Device::CreateCommandAllocator", hr); + } + + hr = device_->CreateCommandList( + 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 = staging_slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close(initial)", hr); + } + + staging_slot.file_read_complete_event.reset(CreateEventExW( + nullptr, nullptr, CREATE_EVENT_MANUAL_RESET, + EVENT_MODIFY_STATE | SYNCHRONIZE)); + if (!staging_slot.file_read_complete_event) { + return Win32Error("CreateEventExW(read)", GetLastError()); + } + } + + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::LoadInternal( + const std::vector& ranges, + BufferCollection& result, + const CancellationToken& cancellation) { + std::lock_guard lock(load_mutex_); + result.Clear(); + + ORT_RETURN_IF( + 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, + "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); + } + + // 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 = + AllocateDestinationGpuBuffers(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; + }); + + // 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 + // every error and cancellation path. + auto cleanup = gsl::finally([&]() noexcept { + if (!load_succeeded) { + DrainActiveReads(); + WaitForFenceUncancelled(last_submitted_fence); + } + if (allocation_future.valid()) { + try { + (void)allocation_future.get(); + } catch (...) { + } + } + }); + + 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(PrepareUploadStagingSlots(maximum_alignment, cancellation)); + + BufferCollection 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) { + 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()) { + return status; + } + } + } + + ORT_RETURN_IF_ERROR(ensure_allocation()); + ORT_RETURN_IF_ERROR( + WaitForFence(last_submitted_fence, cancellation)); + ORT_RETURN_IF_ERROR( + TransitionDestinationBuffersToCommonState(loaded_batch, cancellation, last_submitted_fence)); + + result = std::move(loaded_batch); + load_succeeded = true; + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::AllocateDestinationGpuBuffers( + const std::vector& sizes, + BufferCollection& 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])); + } 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])); + } + if (FAILED(hr)) { + return HResultError( + placement.committed + ? "ID3D12Device::CreateCommittedResource(destination)" + : "ID3D12Device::CreatePlacedResource", + hr); + } + } + + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::CreateFileReadPlan( + 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. + std::map file_indices; + for (size_t range_index = 0; + range_index < ranges.size(); ++range_index) { + const auto& range = ranges[range_index]; + const auto [file_index_it, inserted] = + file_indices.try_emplace(range.path, files.size()); + if (inserted) { + FileReadPlan file; + file.path = range.path; + files.push_back(std::move(file)); + } + files[file_index_it->second].buffer_source_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 Win32FileError( + "CreateFileW(unbuffered)", file.path, 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."); + } + 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 Win32FileError( + "CreateFileW(buffered tail)", file.path, GetLastError()); + } + } + + std::sort( + 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.buffer_source_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; + if (!TryAlignUp( + range.offset + range.length, + 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, file.size)}; + 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.read_regions.push_back(region); + } + } + } + + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::PrepareUploadStagingSlots( + uint64_t alignment, + const CancellationToken& cancellation) { + for (auto& staging_slot : upload_staging_slots_) { + ORT_RETURN_IF_ERROR( + 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."); + } + 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::IssueFileRead( + HANDLE file, + UploadStagingSlot& staging_slot, + uint64_t offset, + DWORD size) { + if (!ResetEvent(staging_slot.file_read_complete_event.get())) { + return Win32Error("ResetEvent(read)", GetLastError()); + } + staging_slot.overlapped = {}; + staging_slot.overlapped.Offset = static_cast(offset); + staging_slot.overlapped.OffsetHigh = + static_cast(offset >> 32); + 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, staging_slot.mapped, size, nullptr, + &staging_slot.overlapped)) { + const DWORD error = GetLastError(); + if (error != ERROR_IO_PENDING) { + staging_slot.read_active = false; + staging_slot.read_file = INVALID_HANDLE_VALUE; + return Win32Error("ReadFile", error); + } + } + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::ProcessFileReadCompletion( + uint64_t file_size, + UploadStagingSlot& staging_slot, + DWORD& bytes_read) { + if (!GetOverlappedResult( + staging_slot.read_file, &staging_slot.overlapped, + &bytes_read, FALSE)) { + return Win32Error( + "GetOverlappedResult", GetLastError()); + } + staging_slot.read_active = false; + staging_slot.read_file = INVALID_HANDLE_VALUE; + const uint64_t required_bytes = + std::min( + staging_slot.requested_read_bytes, + file_size - staging_slot.file_offset); + if (bytes_read < required_bytes) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Read ended before the required file range."); + } + return common::Status::OK(); +} + +template +common::Status D3D12FileBufferLoader::ReadFileRegion( + FileReadPlan& file, + const FileReadRegion& region, + 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 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& staging_slot : upload_staging_slots_) { + if (next_offset >= region.end) { + break; + } + if (staging_slot.read_active) { + continue; + } + if (staging_slot.fence_value != 0 && + copy_fence_->GetCompletedValue() < + staging_slot.fence_value) { + continue; + } + if (cancellation.IsCancellationRequested()) { + return CancelledStatus(); + } + const DWORD bytes = static_cast( + 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( + IssueFileRead( + use_buffered_tail + ? file.buffered_tail_file.get() + : file.file.get(), + staging_slot, + next_offset, bytes)); + next_offset += bytes; + } + return common::Status::OK(); + }; + + ORT_RETURN_IF_ERROR(issue_available()); + DWORD cancellation_poll_milliseconds = + kInitialCancellationPollMilliseconds; + while (true) { + std::array handles{}; + std::array staging_slot_indices{}; + DWORD handle_count = 0; + for (size_t index = 0; + index < upload_staging_slots_.size(); ++index) { + if (upload_staging_slots_[index].read_active) { + handles[handle_count] = + upload_staging_slots_[index].file_read_complete_event.get(); + staging_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& staging_slot : upload_staging_slots_) { + if (!staging_slot.read_active && + staging_slot.fence_value != 0) { + next_fence = + std::min(next_fence, staging_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, + 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 || + wait_result >= WAIT_OBJECT_0 + handle_count) { + return Win32Error( + "WaitForMultipleObjects(read)", + GetLastError()); + } + + auto& staging_slot = + upload_staging_slots_[staging_slot_indices[wait_result - WAIT_OBJECT_0]]; + DWORD bytes_read = 0; + ORT_RETURN_IF_ERROR( + ProcessFileReadCompletion( + file.size, staging_slot, bytes_read)); + + if (!allocation_ready) { + ORT_RETURN_IF_ERROR(ensure_allocation()); + } + ORT_RETURN_IF_ERROR( + SubmitStagingToDestinationCopies( + staging_slot, staging_slot.file_offset, bytes_read, + file.buffer_source_indices, ranges, batch, + last_submitted_fence)); + cancellation_poll_milliseconds = + kInitialCancellationPollMilliseconds; + ORT_RETURN_IF_ERROR(issue_available()); + } + + return common::Status::OK(); +} + +common::Status D3D12FileBufferLoader::SubmitStagingToDestinationCopies( + UploadStagingSlot& staging_slot, + uint64_t chunk_begin, + uint64_t bytes, + const std::vector& buffer_source_indices, + const std::vector& ranges, + BufferCollection& batch, + uint64_t& last_submitted_fence) { + HRESULT hr = staging_slot.allocator->Reset(); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandAllocator::Reset", hr); + } + 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 : buffer_source_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) { + staging_slot.command_list->CopyBufferRegion( + batch.buffers[range_index].Get(), + copy_begin - range.offset, + staging_slot.resource.Get(), + copy_begin - chunk_begin, + copy_end - copy_begin); + } + + // 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}); + staging_slot.command_list->CopyBufferRegion( + batch.buffers[range_index].Get(), + padding_offset, + staging_slot.resource.Get(), + source_offset - chunk_begin, + copy_size); + padding_offset += copy_size; + } else { + ++padding_offset; + } + } + } + + hr = staging_slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close", hr); + } + ID3D12CommandList* command_lists[] = { + staging_slot.command_list.Get()}; + copy_queue_->ExecuteCommandLists(1, command_lists); + + const uint64_t fence_value = ++next_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::TransitionDestinationBuffersToCommonState( + BufferCollection& batch, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence) { + if (batch.buffers.empty()) { + return common::Status::OK(); + } + + auto& staging_slot = upload_staging_slots_.front(); + HRESULT hr = staging_slot.allocator->Reset(); + if (FAILED(hr)) { + return HResultError( + "ID3D12CommandAllocator::Reset(transition)", hr); + } + hr = staging_slot.command_list->Reset( + staging_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.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); + } + staging_slot.command_list->ResourceBarrier( + static_cast(barriers.size()), + barriers.data()); + + hr = staging_slot.command_list->Close(); + if (FAILED(hr)) { + return HResultError( + "ID3D12GraphicsCommandList::Close(transition)", hr); + } + ID3D12CommandList* command_lists[] = { + staging_slot.command_list.Get()}; + copy_queue_->ExecuteCommandLists(1, command_lists); + + const uint64_t fence_value = ++next_fence_value_; + 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, 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, + // 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)) { + return common::Status::OK(); + } + const HRESULT removal_reason = + device_->GetDeviceRemovedReason(); + if (FAILED(removal_reason)) { + return HResultError( + "D3D12 device removed after command submission", + removal_reason); + } + 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); +} + +common::Status D3D12FileBufferLoader::WaitForFence( + uint64_t value, + const CancellationToken& cancellation) { + 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 = + copy_fence_->SetEventOnCompletion( + value, copy_fence_complete_event_.get()); + if (FAILED(hr)) { + return HResultError( + "ID3D12Fence::SetEventOnCompletion", hr); + } + + DWORD cancellation_poll_milliseconds = + kInitialCancellationPollMilliseconds; + while (true) { + const DWORD wait_result = WaitForSingleObject( + copy_fence_complete_event_.get(), + cancellation_poll_milliseconds); + 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(); + } + cancellation_poll_milliseconds = + std::min( + cancellation_poll_milliseconds * 2, + kMaximumCancellationPollMilliseconds); + } +} + +void D3D12FileBufferLoader::WaitForFenceUncancelled( + uint64_t value) noexcept { + if (value == 0 || + copy_fence_->GetCompletedValue() >= value) { + return; + } + if (SUCCEEDED(copy_fence_->SetEventOnCompletion( + value, copy_fence_complete_event_.get())) && + WaitForSingleObject( + copy_fence_complete_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::DrainActiveReads() noexcept { + for (auto& staging_slot : upload_staging_slots_) { + if (!staging_slot.read_active) { + continue; + } + (void)CancelIoEx( + staging_slot.read_file, &staging_slot.overlapped); + } + for (auto& staging_slot : upload_staging_slots_) { + if (!staging_slot.read_active) { + continue; + } + DWORD bytes_read = 0; + (void)GetOverlappedResult( + staging_slot.read_file, &staging_slot.overlapped, + &bytes_read, TRUE); + staging_slot.read_active = false; + staging_slot.read_file = INVALID_HANDLE_VALUE; + } +} + +common::Status D3D12FileBufferLoader::Create( + ID3D12Device* device, + std::unique_ptr& loader, + const Config& config) noexcept { + loader.reset(); + try { + 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( + 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, + BufferCollection& result, + const CancellationToken& cancellation) noexcept { + try { + return LoadInternal(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..961a5ca0c2bf8 --- /dev/null +++ b/onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h @@ -0,0 +1,153 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#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 BufferSource { + 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 BufferCollection { + 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; + + // 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, + BufferCollection& result, + const CancellationToken& cancellation = {}) noexcept; + + private: + struct FileReadRegion; + struct FileReadPlan; + struct UploadStagingSlot; + struct AllocationResult; + + D3D12FileBufferLoader( + ID3D12Device* device, + const Config& config) noexcept; + + common::Status Initialize(); + common::Status LoadInternal( + const std::vector& ranges, + BufferCollection& result, + const CancellationToken& cancellation); + common::Status AllocateDestinationGpuBuffers( + const std::vector& sizes, + BufferCollection& batch); + common::Status CreateFileReadPlan( + const std::vector& ranges, + const CancellationToken& cancellation, + std::vector& files); + common::Status PrepareUploadStagingSlots( + uint64_t alignment, + const CancellationToken& cancellation); + common::Status IssueFileRead( + HANDLE file, + UploadStagingSlot& staging_slot, + uint64_t offset, + DWORD size); + common::Status ProcessFileReadCompletion( + uint64_t file_size, + UploadStagingSlot& staging_slot, + DWORD& bytes_read); + template + common::Status ReadFileRegion( + FileReadPlan& file, + const FileReadRegion& region, + const std::vector& ranges, + BufferCollection& batch, + bool& allocation_ready, + EnsureAllocationFn& ensure_allocation, + const CancellationToken& cancellation, + uint64_t& last_submitted_fence); + common::Status SubmitStagingToDestinationCopies( + UploadStagingSlot& staging_slot, + uint64_t chunk_begin, + uint64_t bytes, + const std::vector& buffer_source_indices, + const std::vector& ranges, + BufferCollection& batch, + uint64_t& last_submitted_fence); + 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, BufferCollection& batch); + 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 copy_fence_complete_event_; + std::vector upload_staging_slots_; + uint64_t next_fence_value_ = 0; + bool upload_error_gpu_completion_unknown_ = false; + std::mutex load_mutex_; +}; + +} // namespace d3d12 +} // namespace windows +} // namespace onnxruntime 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..08d197a485139 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc @@ -0,0 +1,908 @@ +// 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 true; + } + if (other.path < path) { + return false; + } + if (name != other.name) { + return name < other.name; + } + if (offset != other.offset) { + return offset < other.offset; + } + return length < other.length; + } +}; + +struct AcceleratedTensorLoadEntry { + TensorKey key; + ComPtr resource; + std::shared_ptr heap_owner; + wgpu::SharedBufferMemory memory; + wgpu::Buffer buffer; + bool access_started = false; + bool claimed = false; +}; + +} // namespace + +struct D3D12AcceleratedLoadBatch { + 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 load_candidates_committed = false; +}; + +namespace { + +struct D3D12AcceleratedLoadMetrics { + double load_ms = 0.0; + size_t loaded_ranges = 0; +}; + +struct CancellationState { + const std::function* external = nullptr; + const std::atomic* load_end_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.load_end_requested != nullptr && + state.load_end_requested->load(std::memory_order_relaxed)) { + return true; + } + if (state.external == nullptr || !*state.external) { + return false; + } + return InvokeCancellationNoThrow(*state.external); +} + +common::Status LoadBatchToD3D12( + D3D12AcceleratedLoadBatch& batch, + D3D12FileBufferLoader& loader, + const std::function& is_canceled, + const std::atomic& load_end_requested, + D3D12AcceleratedLoadMetrics& metrics) { + if (load_end_requested.load(std::memory_order_relaxed) || + InvokeCancellationNoThrow(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, &load_end_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]; + 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 RegisterLoadCandidateForBatch( + D3D12AcceleratedLoadBatch& 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, + D3D12AcceleratedLoadBatch::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, "\" from \"", key.path.string(), + "\" range [", offset, ", ", offset + length, ")."); + 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( + const WebGpuContext& context) { + ORT_RETURN_IF_NOT( + context.SelectedBackendType() == + 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(); +} + +// Shared by the loader and allocator so imported resources outlive initializer loading. +class D3D12ImportedBufferRegistry { + public: + ~D3D12ImportedBufferRegistry(); + + std::mutex mutex; + IAllocator* allocator = nullptr; + std::unordered_map< + WGPUBuffer, + std::unique_ptr> + imported_allocations; +}; + +D3D12ImportedBufferRegistry:: + ~D3D12ImportedBufferRegistry() { + std::vector> + allocations; + { + std::lock_guard lock{mutex}; + allocations.reserve( + imported_allocations.size()); + for (auto& entry : imported_allocations) { + allocations.push_back(std::move(entry.second)); + } + imported_allocations.clear(); + } + for (auto& allocation : allocations) { + EndAccessNoThrow(*allocation); + } +} + +class D3D12AcceleratedWebGpuAllocator final + : public IAllocator { + public: + D3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + int context_id, + std::function + recording_getter, + std::shared_ptr + buffer_registry) + : IAllocator( + OrtMemoryInfo( + WEBGPU_BUFFER, + OrtAllocatorType::OrtReadOnlyAllocator, + WebGpuDevice(context_id), + OrtMemTypeDefault)), + context_{context}, + recording_getter_{ + std::move(recording_getter)}, + buffer_registry_{std::move(buffer_registry)} { + } + + 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_(); + return context_.InitializerBufferManager().Create( + recording, size, usage); + } + + void Free(void* p) override { + if (p == nullptr) { + return; + } + + auto& recording = recording_getter_(); + std::unique_ptr imported; + { + std::lock_guard lock{ + buffer_registry_->mutex}; + const auto iterator = + buffer_registry_->imported_allocations.find( + static_cast(p)); + if (iterator != + buffer_registry_->imported_allocations.end()) { + imported = std::move(iterator->second); + buffer_registry_->imported_allocations.erase( + iterator); + } + } + + if (imported) { + if (recording.has_unsubmitted_work.load( + std::memory_order_acquire)) { + std::shared_ptr deferred{ + imported.release(), + [](ImportedAllocation* allocation) { + EndAccessNoThrow(*allocation); + delete allocation; + }}; + { + 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); + } + return; + } + + context_.InitializerBufferManager().Release( + static_cast(p), &recording); + } + + private: + WebGpuContext& context_; + std::function + recording_getter_; + std::shared_ptr + buffer_registry_; +}; + +AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + int context_id, + std::function + recording_getter, + std::shared_ptr& + out_buffer_registry) { + auto buffer_registry = + std::make_shared(); + auto allocator = + std::make_shared< + D3D12AcceleratedWebGpuAllocator>( + context, context_id, 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(device_support_resolution_once_, [this]() { + common::Status support_status = + common::Status::OK(); + if (IsWeightLoadAccelerationEnabled(acceleration_mode_)) { + support_status = + CheckD3D12AcceleratedExternalWeightsSupport( + context_); + } + device_support_resolution_status_ = + ResolveWeightLoadAccelerationMode( + 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. " + "Reason: " + << support_status.ErrorMessage(); + } + }); +} + +common::Status D3D12AcceleratedExternalDataLoader::EnsureFileLoader() const { + if (file_to_buffer_loader_) { + return common::Status::OK(); + } + return D3D12FileBufferLoader::Create( + context_.WeightLoadingD3D12Device(), file_to_buffer_loader_); +} + +D3D12AcceleratedExternalDataLoader:: + D3D12AcceleratedExternalDataLoader( + WebGpuContext& context, + int context_id, + std::shared_ptr< + D3D12ImportedBufferRegistry> + buffer_registry, + WeightLoadAccelerationMode 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."); +} + +D3D12AcceleratedExternalDataLoader:: + ~D3D12AcceleratedExternalDataLoader() { + EndLoad(); +} + +bool D3D12AcceleratedExternalDataLoader::CanLoad( + const OrtMemoryInfo& target_memory_info) const { + ResolveSupport(); + return target_memory_info.device == WebGpuDevice(context_id_) && + target_memory_info.name == WEBGPU_BUFFER && + resolved_acceleration_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 { + ResolveSupport(); + return target_device == WebGpuDevice(context_id_) && + (resolved_acceleration_enabled_ || + IsWeightLoadAccelerationRequired(acceleration_mode_)); +} + +common::Status +D3D12AcceleratedExternalDataLoader::BeginLoad() const { + EndLoad(); + ResolveSupport(); + ORT_RETURN_IF_ERROR(device_support_resolution_status_); + if (!resolved_acceleration_enabled_) { + return common::Status::OK(); + } + load_end_requested_.store( + false, std::memory_order_relaxed); + accelerated_load_batch_ = + std::make_unique(); + return common::Status::OK(); +} + +common::Status +D3D12AcceleratedExternalDataLoader::RegisterLoadCandidate( + 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( + accelerated_load_batch_ != nullptr && + !accelerated_load_batch_->load_candidates_committed, + "D3D12 accelerated initializer batch has not been started."); + return RegisterLoadCandidateForBatch( + *accelerated_load_batch_, env, data_file_path, + tensor_name, data_offset, data_length); +} + +common::Status +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_->load_candidates_committed, + "D3D12 accelerated initializer batch has not been started."); + auto& batch = *accelerated_load_batch_; + const auto fail_or_fallback = + [this, &is_canceled]( + const common::Status& status) + -> common::Status { + if (status.Code() == + common::MODEL_LOAD_CANCELED || + InvokeCancellationNoThrow(is_canceled)) { + return ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled."); + } + if (IsWeightLoadAccelerationRequired( + acceleration_mode_)) { + return status; + } + LOGS_DEFAULT(WARNING) + << "D3D12 accelerated initializer loading failed; " + "using the ordinary WebGPU initializer path: " + << status.ErrorMessage(); + resolved_acceleration_enabled_ = false; + EndLoad(); + return common::Status::OK(); + }; + + if (batch.load_range_count == 0) { + if (InvokeCancellationNoThrow(is_canceled)) { + return fail_or_fallback(ORT_MAKE_STATUS( + ONNXRUNTIME, MODEL_LOAD_CANCELED, + "D3D12 accelerated initializer loading was canceled.")); + } + batch.load_candidates_committed = true; + return common::Status::OK(); + } + + D3D12AcceleratedLoadMetrics load_metrics; + const auto loader_status = EnsureFileLoader(); + if (!loader_status.IsOK()) { + return fail_or_fallback(loader_status); + } + const auto load_status = LoadBatchToD3D12( + batch, *file_to_buffer_loader_, is_canceled, + load_end_requested_, load_metrics); + if (!load_status.IsOK()) { + return fail_or_fallback(load_status); + } + + if (InvokeCancellationNoThrow(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( + context_.D3D12SharedResourceFeaturesAvailable(), + "D3D12 accelerated external weights require Dawn " + "SharedBufferMemoryD3D12Resource and SharedFenceDXGISharedHandle features."); + for (auto& tensor : batch.tensors) { + if (InvokeCancellationNoThrow(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 = + 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(); + // 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 | + 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.load_candidates_committed = true; + + LOGS_DEFAULT(VERBOSE) + << "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::EndLoad() + const noexcept { + load_end_requested_.store( + true, std::memory_order_relaxed); + + ORT_TRY { + if (accelerated_load_batch_) { + for (auto& tensor : + accelerated_load_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; + } + } + accelerated_load_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( + accelerated_load_batch_ != nullptr && + 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."); + ORT_RETURN_IF_NOT( + allocator != nullptr, + "D3D12 accelerated initializer requires its device allocator."); + ORT_RETURN_IF_NOT( + allocator.get() == + buffer_registry_->allocator, + "D3D12 accelerated initializer was passed a different allocator."); + + const size_t length = + static_cast(data_length); + const auto file_iterator = + accelerated_load_batch_->files.find( + data_file_path.lexically_normal()); + ORT_RETURN_IF( + file_iterator == accelerated_load_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 = + accelerated_load_batch_->tensors_by_key.find(key); + ORT_RETURN_IF( + iterator == + accelerated_load_batch_->tensors_by_key.end(), + "No prepared D3D12 accelerated initializer matches \"", + tensor_name, + "\" and the requested file range."); + + auto& prepared = + accelerated_load_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{ + buffer_registry_->mutex}; + ORT_RETURN_IF( + buffer_registry_->imported_allocations + .find(buffer) != + buffer_registry_->imported_allocations + .end(), + "D3D12 accelerated imported buffer was registered twice."); + buffer_registry_->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..13a5969ff4866 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/d3d12_external_data_loader.h @@ -0,0 +1,93 @@ +// 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 + +#include "core/framework/allocator.h" +#include "core/framework/external_data_loader.h" +#include "core/providers/webgpu/webgpu_provider_options.h" + +namespace onnxruntime { +namespace windows { +namespace d3d12 { +class D3D12FileBufferLoader; +} // namespace d3d12 +} // namespace windows + +namespace webgpu { + +class WebGpuContext; +struct CommandRecordingState; + +common::Status CheckD3D12AcceleratedExternalWeightsSupport(const WebGpuContext& context); + +common::Status ResolveWeightLoadAccelerationMode( + WeightLoadAccelerationMode mode, + const common::Status& support_status, + bool& enabled); + +class D3D12ImportedBufferRegistry; +struct D3D12AcceleratedLoadBatch; + +AllocatorPtr CreateD3D12AcceleratedWebGpuAllocator( + WebGpuContext& context, + int context_id, + std::function recording_getter, + std::shared_ptr& out_buffer_registry); + +class D3D12AcceleratedExternalDataLoader final : public IExternalDataLoader { + public: + D3D12AcceleratedExternalDataLoader( + WebGpuContext& context, + int context_id, + std::shared_ptr buffer_registry, + 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 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, + FileOffsetType data_offset, + SafeInt data_length, + const std::shared_ptr& allocator, + Tensor& tensor) const override; + + private: + void ResolveSupport() const; + 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_; + 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 load_end_requested_{false}; + mutable std::unique_ptr file_to_buffer_loader_; + mutable std::unique_ptr accelerated_load_batch_; +}; + +} // namespace webgpu +} // namespace onnxruntime + +#endif // defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index cc5567b9ce6bf..2a7c4da531e4f 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -3,11 +3,23 @@ #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" +#if !defined(USE_EXTERNAL_DAWN) +#include "dawn/native/D3D12Backend.h" +#endif +#endif + #if defined(__GNUC__) #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wstrict-aliasing" @@ -42,8 +54,24 @@ namespace onnxruntime { namespace webgpu { +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +namespace { + +struct RequestAdapterOptionsLuid : wgpu::ChainedStruct { + RequestAdapterOptionsLuid() { + sType = static_cast(WGPUSType_RequestAdapterOptionsLUID); + } + + LUID adapter_luid{}; +}; + +} // namespace +#endif + +WebGpuContext::~WebGpuContext() = default; + void WebGpuContext::Initialize(const WebGpuContextConfig& config) { - std::call_once(init_flag_, [this, &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; @@ -52,11 +80,11 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { 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; // 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 @@ -66,8 +94,6 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { // 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; } @@ -88,6 +114,105 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { 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_) && + static_cast(config.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(); + 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) @@ -111,27 +236,54 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { 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)); + 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 = {}; @@ -219,6 +371,24 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { 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; + } +#endif LOGS_DEFAULT(VERBOSE) << "WebGPU EP Context is created for: Instance=" << instance_.Get() << ", Device=" << device_.Get() << "."; // cache device queue @@ -246,22 +416,25 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { for (size_t i = 0; i < supported_features.featureCount; i++) { device_features_.insert(supported_features.features[i]); } + // cache adapter info #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. + // 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."); -#endif - // cache adapter info if (DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) { adapter_info_.nextInChain = &subgroup_matrix_configs_; } +#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. + selected_backend_type_ = adapter_info_.backendType; +#endif // create buffer manager buffer_mgr_ = BufferManagerFactory::Create(*this, @@ -269,6 +442,8 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { 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, @@ -294,6 +469,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) && + selected_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 " @@ -331,6 +516,12 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { } } +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) +ID3D12Device* WebGpuContext::WeightLoadingD3D12Device() const { + 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) { @@ -926,6 +1117,27 @@ 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, + }; + // 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 return required_features; } @@ -1255,6 +1467,17 @@ 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) + std::vector> pending_release_callbacks; + { + std::lock_guard callbacks_lock{ + recording.pending_release_callbacks_mutex}; + pending_release_callbacks.swap(recording.pending_release_callbacks); + } + for (auto& release : pending_release_callbacks) { + release(); + } +#endif recording.command_encoder = nullptr; recording.num_pending_dispatches = 0; return status; @@ -1442,9 +1665,8 @@ WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& co } if (context_id == 0) { - // context ID is preserved for the default context. User cannot use context ID 0 as a custom context. 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 { diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 2fd9e77871e24..e02777b1d78b3 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,20 @@ #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 { + 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 +129,10 @@ 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::mutex pending_release_callbacks_mutex; + 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_; } @@ -322,6 +337,15 @@ 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; } +#if defined(_WIN32) && defined(ENABLE_D3D12_FILE_LOADING) + ID3D12Device* WeightLoadingD3D12Device() const; + bool D3D12SharedResourceFeaturesAvailable() const { + return d3d12_shared_resource_features_available_; + } + wgpu::BackendType SelectedBackendType() const { + return selected_backend_type_; + } +#endif // // Get Split-K configuration. @@ -492,6 +516,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 selected_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..ec2ed34439ab6 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 @@ -617,12 +620,40 @@ 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) && + context_.HasDevice()) { + accelerated_initializer_allocator_ = + CreateD3D12AcceleratedWebGpuAllocator( + context_, + context_id_, + [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); @@ -642,12 +673,18 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id, std::vector WebGpuExecutionProvider::CreatePreferredAllocators() { const bool device_free = !context_.HasDevice(); return { - // allocator for initializers - CreateWebGpuAllocator( - context_id_, - device_free, - [this]() -> const webgpu::BufferManager& { return InitializerBufferManager(); }, - [this]() -> webgpu::CommandRecordingState& { return Recording(); }, true), + // 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), // default allocator CreateWebGpuAllocator( context_id_, @@ -770,6 +807,16 @@ std::unique_ptr WebGpuExecutionProvider::GetDataTran 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_, context_id_, 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..d1009143bb18a 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 D3D12ImportedBufferRegistry; +#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..843bb5cd8893b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -47,7 +47,61 @@ struct WebGpuProviderFactory : IExecutionProviderFactory { namespace { -WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) { +constexpr const char* WeightLoadAccelerationModeName( + WeightLoadAccelerationMode mode) { + switch (mode) { + case WeightLoadAccelerationMode::Off: + return kWeightLoadAcceleration_Off; + case WeightLoadAccelerationMode::Preferred: + return kWeightLoadAcceleration_Preferred; + case WeightLoadAccelerationMode::Required: + return kWeightLoadAcceleration_Required; + } + return "unknown"; +} + +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\"."); +} + +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; @@ -72,6 +126,9 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) } } + webgpu_ep_config.weight_load_acceleration_mode = + weight_load_acceleration_mode; + if (std::string pool_generations_str; config_options.TryGetConfigEntry(kSessionBufferPoolGenerations, pool_generations_str)) { size_t pool_generations = 0; @@ -167,6 +224,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: " + << 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; @@ -176,9 +235,14 @@ 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 = + weight_load_acceleration_mode; + 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(), @@ -405,11 +469,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/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/test/framework/external_data_loader_test.cc b/onnxruntime/test/framework/external_data_loader_test.cc index 35ecd44477afd..43b1ff1dbd3b3 100644 --- a/onnxruntime/test/framework/external_data_loader_test.cc +++ b/onnxruntime/test/framework/external_data_loader_test.cc @@ -1,11 +1,19 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include "core/framework/external_data_loader_manager.h" + +#include +#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,104 @@ TEST_F(ExternalDataLoaderLifetimeTest, RecreatesLoaderAfterFactoryFailure) { } #endif +#endif + +#if defined(ENABLE_D3D12_FILE_LOADING) +class BatchLifecycleExternalDataLoader final : public IExternalDataLoader { + public: + enum class FailurePoint { + None, + Begin, + CommitCandidates, + }; + + 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 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 EndLoad() const noexcept override { + ++end_count; + } + + mutable int begin_count = 0; + mutable int commit_candidates_count = 0; + mutable int end_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.CommitLoadCandidates([]() { return false; })); + manager.EndLoad(); +} + +TEST(ExternalDataLoaderManagerTest, BeginFailureEndsEveryLoader) { + 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->end_count, 1); + EXPECT_EQ(failing_ptr->end_count, 1); +} + +TEST(ExternalDataLoaderManagerTest, CommitCandidatesFailureEndsEveryLoader) { + ExternalDataLoaderManager manager; + auto failing = std::make_unique( + BatchLifecycleExternalDataLoader::FailurePoint::CommitCandidates); + auto* failing_ptr = failing.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(uncommitted))); + + ASSERT_STATUS_OK(manager.BeginLoad()); + 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 + } // 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..7109b957a873b 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")}; @@ -76,6 +175,328 @@ 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: 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.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("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(); +} + // 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 @@ -200,5 +621,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 0b9473ae00257..7a09b4eead7b5 100644 --- a/onnxruntime/test/providers/webgpu/webgpu_context_test.cc +++ b/onnxruntime/test/providers/webgpu/webgpu_context_test.cc @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -20,6 +23,8 @@ #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" @@ -27,7 +32,17 @@ #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/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h" +#include "core/providers/webgpu/d3d12_external_data_loader.h" +#endif #if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) #include "dawn/native/DawnNative.h" @@ -103,6 +118,25 @@ 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; +} + +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()); @@ -120,9 +154,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 +178,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()); @@ -1119,6 +1155,858 @@ 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, RequiredWeightLoadAccelerationDoesNotExposeLoaderWithoutDevice) { + 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); +#else + EXPECT_THROW(factory->CreateProvider(), OnnxRuntimeException); +#endif + }, + 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, + 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_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); + + 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].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( + 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(); + } + EXPECT_TRUE(context.DeviceHasFeature( + wgpu::FeatureName::SharedBufferMemoryD3D12Resource)); + EXPECT_TRUE(context.DeviceHasFeature( + wgpu::FeatureName::SharedFenceDXGISharedHandle)); + + auto loader = required_ep->GetExternalDataLoader(); + ASSERT_NE(loader, nullptr); + EXPECT_STATUS_OK(loader->BeginLoad()); + loader->EndLoad(); + }); +} + +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); + if (!support_status.IsOK()) { + GTEST_SKIP() << support_status.ErrorMessage(); + } + + 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::BufferCollection batch; + ASSERT_STATUS_OK(loader->Load( + {{data_path.native(), 0, kDataSize}}, batch)); + ASSERT_EQ(batch.buffers.size(), 1u); + + auto* device = context.WeightLoadingD3D12Device(); + ExpectD3D12BufferContents(device, batch, {expected}); +} + +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"); + + 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()); + } + + 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(); + } + + 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::BufferCollection 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) { + 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_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()); + const auto& allocator = allocators.front(); + ASSERT_STATUS_OK(loader->BeginLoad()); + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( + Env::Default(), data_path, "weights", 0, sizeof(expected))); + ASSERT_STATUS_OK(loader->CommitLoadCandidates([]() { return false; })); + + auto& webgpu_ep = *static_cast(ep.get()); + auto& recording = webgpu_ep.Recording(); + 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(); + 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) { + 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->RegisterLoadCandidate( + Env::Default(), data_path, "weights", kDataOffset, sizeof(expected))); + 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->RegisterLoadCandidate( + Env::Default(), data_path, "separated", kSeparatedDataOffset, + sizeof(separated_expected))); + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( + Env::Default(), other_data_path, "other_weights", kDataOffset, + sizeof(other_expected))); + ASSERT_STATUS_OK(loader->RegisterLoadCandidate( + Env::Default(), data_path, "empty", + kSeparatedDataOffset + sizeof(separated_expected), 0)); + ASSERT_STATUS_OK(loader->CommitLoadCandidates([]() { 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; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsDisableCPUEPFallback, "1")); + 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)); + 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), + 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_Required)) + ->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->RegisterLoadCandidate( + Env::Default(), data_path, "weights", 0, sizeof(data))); + const auto status = loader->CommitLoadCandidates([]() { return true; }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->EndLoad(); +} + +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->CommitLoadCandidates([]() { return true; }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->EndLoad(); +} + +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->RegisterLoadCandidate( + Env::Default(), data_path, "weights", 0, sizeof(data))); + + size_t cancellation_checks = 0; + const auto status = loader->CommitLoadCandidates([&cancellation_checks]() { + return cancellation_checks++ != 0; + }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->EndLoad(); +} + +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->RegisterLoadCandidate( + Env::Default(), data_path, "weights", 0, sizeof(data))); + ASSERT_TRUE(std::filesystem::remove(data_path)); + EXPECT_STATUS_OK(loader->CommitLoadCandidates([]() { 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->RegisterLoadCandidate( + Env::Default(), data_path, "weights", 0, sizeof(data))); + ASSERT_TRUE(std::filesystem::remove(data_path)); + const auto status = loader->CommitLoadCandidates([]() { return true; }); + EXPECT_EQ(status.Code(), common::MODEL_LOAD_CANCELED); + loader->EndLoad(); +} + +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."; 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) {