From eeb8fc093191d8cc29aac5369f07c44d4779a868 Mon Sep 17 00:00:00 2001
From: Wanming Lin <wanming.lin@intel.com>
Date: Tue, 9 Jul 2024 23:49:58 +0800
Subject: [PATCH] [WebNN EP] Release WebNN MLGraphBuilder after Compile to free
 memory (#21200)

This would help release the constants bound by the MLGraphBuilder.
---
 .../core/providers/webnn/webnn_execution_provider.cc   | 10 ++++++++++
 .../core/providers/webnn/webnn_execution_provider.h    |  4 ++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/core/providers/webnn/webnn_execution_provider.cc b/onnxruntime/core/providers/webnn/webnn_execution_provider.cc
index 3f8f17048b..e839d6d17b 100644
--- a/onnxruntime/core/providers/webnn/webnn_execution_provider.cc
+++ b/onnxruntime/core/providers/webnn/webnn_execution_provider.cc
@@ -81,6 +81,13 @@ WebNNExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_view
 
   const auto& logger = *GetLogger();
 
+  if (!wnn_builder_.as<bool>()) {
+    // The GetCapability function may be called again after Compile due to the logic in the
+    // PartitionOnnxFormatModel function (see onnxruntime/core/framework/graph_partitioner.cc).
+    // We need to re-create the wnn_builder_ here to avoid it's been released in last Compile.
+    wnn_builder_ = emscripten::val::global("MLGraphBuilder").new_(wnn_context_);
+  }
+
   const auto node_groups = webnn::GetSupportedNodes(graph_viewer, wnn_builder_, wnn_device_type_, logger);
 
   if (node_groups.empty()) {
@@ -322,6 +329,9 @@ common::Status WebNNExecutionProvider::Compile(const std::vector<FusedNodeAndGra
     node_compute_funcs.push_back(compute_info);
   }
 
+  // Explictly release the WebNN builder to free memory.
+  wnn_builder_ = emscripten::val::undefined();
+
   return Status::OK();
 }
 
diff --git a/onnxruntime/core/providers/webnn/webnn_execution_provider.h b/onnxruntime/core/providers/webnn/webnn_execution_provider.h
index ec02f249b6..83d530efed 100644
--- a/onnxruntime/core/providers/webnn/webnn_execution_provider.h
+++ b/onnxruntime/core/providers/webnn/webnn_execution_provider.h
@@ -42,8 +42,8 @@ class WebNNExecutionProvider : public IExecutionProvider {
   std::shared_ptr<KernelRegistry> GetKernelRegistry() const override;
 
  private:
-  emscripten::val wnn_context_ = emscripten::val::object();
-  emscripten::val wnn_builder_ = emscripten::val::object();
+  emscripten::val wnn_context_ = emscripten::val::undefined();
+  mutable emscripten::val wnn_builder_ = emscripten::val::undefined();
 
   DataLayout preferred_layout_;
   webnn::WebnnDeviceType wnn_device_type_;