Fix WebGPU EP crash on exit

Author: fs-eire

onnxruntime/core/providers/webgpu/webgpu_context.cc

@@ -924,10 +924,12 @@ void WebGpuContext::ReleaseGraphResources(std::vector<webgpu::CapturedCommandInf
   }
 }
 
-std::unordered_map<int32_t, WebGpuContextFactory::WebGpuContextInfo>* WebGpuContextFactory::contexts_ = nullptr;
 std::mutex WebGpuContextFactory::mutex_;
 std::once_flag WebGpuContextFactory::init_default_flag_;
-wgpu::Instance WebGpuContextFactory::default_instance_;
+
+std::unordered_map<int32_t, WebGpuContextFactory::WebGpuContextInfo>* WebGpuContextFactory::contexts_ = nullptr;
+WGPUInstance WebGpuContextFactory::default_instance_ = nullptr;
+bool WebGpuContextFactory::modules_dxc_loaded_ = false;
 
 WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& config) {
   const int context_id = config.context_id;
@@ -960,28 +962,60 @@ WebGpuContext& WebGpuContextFactory::CreateContext(const WebGpuContextConfig& co
 
   std::lock_guard<std::mutex> lock(mutex_);
 
-  // Lazy-allocate the contexts map on first use (heap-allocated to avoid static destruction crash).
-  if (contexts_ == nullptr) {
-    contexts_ = new std::unordered_map<int32_t, WebGpuContextInfo>();
-  }
-
   if (default_instance_ == nullptr) {
     // Create wgpu::Instance
     wgpu::InstanceFeatureName required_instance_features[] = {wgpu::InstanceFeatureName::TimedWaitAny};
     wgpu::InstanceDescriptor instance_desc{};
     instance_desc.requiredFeatures = required_instance_features;
     instance_desc.requiredFeatureCount = sizeof(required_instance_features) / sizeof(required_instance_features[0]);
-    default_instance_ = wgpu::CreateInstance(&instance_desc);
+    default_instance_ = wgpuCreateInstance(&instance_desc);
 
     ORT_ENFORCE(default_instance_ != nullptr, "Failed to create wgpu::Instance.");
+
+    // Lazy-allocate the contexts map on first use (heap-allocated to avoid static destruction crash).
+    if (contexts_ == nullptr) {
+      contexts_ = new std::unordered_map<int32_t, WebGpuContextInfo>();
+    }
+
+// If we are on Windows and the build does not use external Dawn, for D3D12 backend dxil.dll and dxcompiler.dll are required.
+//
+// Dawn will try to load them later, but if the DLLs are loaded by Dawn, it could cause them to be unloaded earlier
+// than the descruction of WebGpuContextFactory, which will cause crash because the resource release can potentially
+// call into dxcompiler.dll. By loading them here, we can make sure they are not unloaded too early.
+#if !defined(__wasm__) && defined(_WIN32) && !defined(USE_EXTERNAL_DAWN)
+    if (config.backend_type == static_cast<int>(WebGpuBackendType::D3D12) && !modules_dxc_loaded_) {
+      auto runtime_path = Env::Default().GetRuntimePath();
+      if (!runtime_path.empty()) {
+        if (modules_ == nullptr) {
+          modules_ = new LibraryHandles();
+        }
+
+        Status status;
+        void* module_handle = nullptr;
+
+        PathString dxil_path = runtime_path + ToPathString(L"dxil.dll");
+        status = Env::Default().LoadDynamicLibrary(dxil_path, false, &module_handle);
+        if (status.IsOK() && module_handle != nullptr) {
+          modules_->Add(dxil_path, module_handle);
+        }
+
+        PathString dxcompiler_path = runtime_path + ToPathString(L"dxcompiler.dll");
+        status = Env::Default().LoadDynamicLibrary(dxcompiler_path, false, &module_handle);
+        if (status.IsOK() && module_handle != nullptr) {
+          modules_->Add(dxcompiler_path, module_handle);
+        }
+        modules_dxc_loaded_ = true;
+      }
+    }
+#endif
   }
 
   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.");
 
-    instance = default_instance_.Get();
+    instance = default_instance_;
   } else {
     // for context ID > 0, user must provide custom WebGPU instance and device.
     ORT_ENFORCE(instance != nullptr && device != nullptr,
@@ -1034,9 +1068,18 @@ void WebGpuContextFactory::ReleaseContext(int context_id) {
 
 void WebGpuContextFactory::Cleanup() {
   std::lock_guard<std::mutex> lock(mutex_);
-  delete contexts_;
-  contexts_ = nullptr;
-  default_instance_ = nullptr;
+  if (contexts_ != nullptr) {
+    delete contexts_;
+    contexts_ = nullptr;
+  }
+  if (default_instance_ != nullptr) {
+    wgpuReleaseInstance(default_instance_);
+    default_instance_ = nullptr;
+  }
+  if (modules_ != nullptr) {
+    delete modules_;
+    modules_ = nullptr;
+  }
 }
 
 WebGpuContext& WebGpuContextFactory::DefaultContext() {

onnxruntime/core/providers/webgpu/webgpu_context.h

@@ -10,6 +10,7 @@
 #include "core/providers/webgpu/webgpu_external_header.h"
 
 #include "core/common/common.h"
+#include "core/framework/library_handles.h"
 #include "core/providers/webgpu/buffer_manager.h"
 #include "core/providers/webgpu/program_manager.h"
 #include "core/providers/webgpu/webgpu_utils.h"
@@ -147,6 +148,9 @@ class WebGpuContextFactory {
  private:
   WebGpuContextFactory() {}
 
+  static std::mutex mutex_;
+  static std::once_flag init_default_flag_;
+
   // Use pointers to heap-allocated objects so that their destructors do NOT run
   // during static destruction at process exit. This avoids crashes when dependent
   // DLLs (e.g. dxcompiler.dll) have already been unloaded by the OS.
@@ -155,9 +159,12 @@ class WebGpuContextFactory {
   // it is reached from OrtEnv::~OrtEnv via CleanupWebGpuContexts().
   // On abnormal/process termination they simply leak, which is safe.
   static std::unordered_map<int32_t, WebGpuContextInfo>* contexts_;
-  static std::mutex mutex_;
-  static std::once_flag init_default_flag_;
-  static wgpu::Instance default_instance_;
+  static WGPUInstance default_instance_;
+
+  // Use a module manager to ensure that dependent DLLs (e.g. dxcompiler.dll) are not unloaded while WebGPU contexts are
+  // still alive, which would cause crashes if the contexts try to call into those DLLs during their destruction.
+  static LibraryHandles* modules_;
+  static bool modules_dxc_loaded_;
 };
 
 // Class WebGpuContext includes all necessary resources for the context.