From a05203b91c96ca5ca0112f9403267c04da55ec78 Mon Sep 17 00:00:00 2001
From: Peishen Yan <peishen.yan@intel.com>
Date: Sat, 18 Jan 2025 09:18:38 +0800
Subject: [PATCH] [WebNN EP] Fix AddInitializersToSkip issues (#23354)

### Description
<!-- Describe your changes. -->
When the onnx model reuses initializers in more than one ops, if one of
the ops wants to add this initializer to the skipped list, but the other
ops still need this initializer, it will cause the process to crash.

Therefore, like other EPs, we count `initializer_usage_`, the number of
occurrences of each initializer in all ops and modify the
`AddInitializersToSkip` to minus the corresponding initializers'
statistic one when adding the specific operators.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
---
 .../providers/webnn/builders/model_builder.cc | 26 ++++++++++++++-----
 .../providers/webnn/builders/model_builder.h  |  2 +-
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.cc b/onnxruntime/core/providers/webnn/builders/model_builder.cc
index e8f116d390..097dd16307 100644
--- a/onnxruntime/core/providers/webnn/builders/model_builder.cc
+++ b/onnxruntime/core/providers/webnn/builders/model_builder.cc
@@ -75,9 +75,18 @@ InitializedTensorSet ModelBuilder::GetInitializerTensors() {
 }
 
 void ModelBuilder::PreprocessInitializers() {
+  const auto& initializers = graph_viewer_.GetAllInitializedTensors();
   const auto& node_indices = graph_viewer_.GetNodesInTopologicalOrder();
   for (size_t i = 0; i < node_indices.size(); i++) {
     const auto* node(graph_viewer_.GetNode(node_indices[i]));
+
+    // find all initializers consumed. AddInitializersToSkip will potentially decrement the usage count.
+    for (const auto* input : node->InputDefs()) {
+      if (input->Exists() && Contains(initializers, input->Name())) {
+        initializer_usage_[input->Name()]++;
+      }
+    }
+
     if (const auto* op_builder = GetOpBuilder(*node)) {
       op_builder->AddInitializersToSkip(*this, *node);
     }
@@ -90,12 +99,11 @@ Status ModelBuilder::RegisterInitializers() {
     const auto& name = tensor.name();
     const auto& shape = tensor.dims();
 
-    // Ignore the following tensors:
-    // 1. Empty tensors: optional tensors can be indicated by an empty name.
-    // 2. Tensors in skipped_initializers_: These are tensors that are not used as WebNN Constants.
-    //    Note: Scalar tensors are excluded because ONNX Runtime will optimize same scalar initializers into one.
-    if (name.empty() || (Contains(skipped_initializers_, name) && !shape.empty()))
+    // skip initializer if there is no remaining usage
+    auto usage_count = initializer_usage_[name];
+    if (usage_count == 0) {
       continue;
+    }
 
     std::vector<int32_t> dims;
     // When the shape is empty, it is scalar initializer that dims = {};
@@ -385,7 +393,13 @@ void ModelBuilder::AddOperand(const std::string& name, const emscripten::val& op
 }
 
 void ModelBuilder::AddInitializerToSkip(const std::string& tensor_name) {
-  skipped_initializers_.insert(tensor_name);
+  // Decrement usage count if this is a known initializer.
+  // For simplicity the OpBuilder::AddInitializersToSkip implementations may call this for arbitrary input names
+  // without first checking if the value is an initializer.
+  auto entry = initializer_usage_.find(tensor_name);
+  if (entry != initializer_usage_.end()) {
+    --entry->second;
+  }
 }
 
 void ModelBuilder::AddInputToSkip(const std::string& input_name) {
diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.h b/onnxruntime/core/providers/webnn/builders/model_builder.h
index 0fc2fa2067..4e2d84f481 100644
--- a/onnxruntime/core/providers/webnn/builders/model_builder.h
+++ b/onnxruntime/core/providers/webnn/builders/model_builder.h
@@ -81,7 +81,7 @@ class ModelBuilder {
 
   InlinedHashMap<std::string, OnnxTensorInfo> input_output_info_;
 
-  InlinedHashSet<std::string> skipped_initializers_;
+  std::unordered_map<std::string, int> initializer_usage_;
   InlinedHashSet<std::string> skipped_inputs_;
 
   uint32_t name_token_{0};