[WebGPU] Accelerated model loading using Win32 and D3D12

.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
           }

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)

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()

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.
     #

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.

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" ||

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<NonNullable<InferenceSession.WebGpuExecutionProviderOption['weightLoadAcceleration']>> =
+      ['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);

onnxruntime/core/framework/external_data_loader.h

@@ -4,16 +4,20 @@
 #pragma once
 
 #include <functional>
-#include <vector>
 #include <filesystem>
+#include <memory>
+#include <string_view>
+#include <vector>
 
 #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<size_t> 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<bool()>& 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<size_t> data_length,
+                                    [[maybe_unused]] const std::shared_ptr<IAllocator>& 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,

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<bool()>& 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

onnxruntime/core/framework/external_data_loader_manager.h

@@ -3,6 +3,8 @@
 
 #pragma once
 
+#include <functional>
+
 #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<bool()>& is_canceled) const;
+  void EndLoad() const noexcept;
+#endif
+
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExternalDataLoaderManager);