resolve comments: add a cleanup function to call in OrtEnv destructor

microsoft · Oct 26, 2024 · 2239f56 · 2239f56
1 parent c86b1cc
commit 2239f56
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc
@@ -674,5 +674,14 @@ WebGpuContext& WebGpuContextFactory::GetContext(int context_id) {
   return *it->second;
 }
 
+void WebGpuContextFactory::Cleanup() {
+  std::lock_guard<std::mutex> lock(mutex_);
+  contexts_.clear();
+}
+
+void CleanupWebGpuContexts() {
+  WebGpuContextFactory::Cleanup();
+}
+
 }  // namespace webgpu
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h
@@ -34,6 +34,8 @@ class WebGpuContextFactory {
                                       ValidationMode validation_mode);
   static WebGpuContext& GetContext(int context_id);
 
+  static void Cleanup();
+
  private:
   WebGpuContextFactory() {}
 

diff --git a/onnxruntime/core/session/ort_env.cc b/onnxruntime/core/session/ort_env.cc
@@ -17,6 +17,14 @@
 using namespace onnxruntime;
 using namespace onnxruntime::logging;
 
+#ifdef USE_WEBGPU
+namespace onnxruntime {
+namespace webgpu {
+void CleanupWebGpuContexts();
+}  // namespace webgpu
+}  // namespace onnxruntime
+#endif
+
 std::unique_ptr<OrtEnv> OrtEnv::p_instance_;
 int OrtEnv::ref_count_ = 0;
 std::mutex OrtEnv::m_;
@@ -26,6 +34,10 @@ OrtEnv::OrtEnv(std::unique_ptr<onnxruntime::Environment> value1)
 }
 
 OrtEnv::~OrtEnv() {
+#ifdef USE_WEBGPU
+  webgpu::CleanupWebGpuContexts();
+#endif
+
 // We don't support any shared providers in the minimal build yet
 #if !defined(ORT_MINIMAL_BUILD)
   UnloadSharedProviders();