Skip to content

Commit d633e57

Browse files
peishenyanashrit-ms
authored andcommitted
[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. -->
1 parent e988ef0 commit d633e57

File tree

2 files changed

+21
-7
lines changed

2 files changed

+21
-7
lines changed

onnxruntime/core/providers/webnn/builders/model_builder.cc

+20-6
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,18 @@ InitializedTensorSet ModelBuilder::GetInitializerTensors() {
7575
}
7676

7777
void ModelBuilder::PreprocessInitializers() {
78+
const auto& initializers = graph_viewer_.GetAllInitializedTensors();
7879
const auto& node_indices = graph_viewer_.GetNodesInTopologicalOrder();
7980
for (size_t i = 0; i < node_indices.size(); i++) {
8081
const auto* node(graph_viewer_.GetNode(node_indices[i]));
82+
83+
// find all initializers consumed. AddInitializersToSkip will potentially decrement the usage count.
84+
for (const auto* input : node->InputDefs()) {
85+
if (input->Exists() && Contains(initializers, input->Name())) {
86+
initializer_usage_[input->Name()]++;
87+
}
88+
}
89+
8190
if (const auto* op_builder = GetOpBuilder(*node)) {
8291
op_builder->AddInitializersToSkip(*this, *node);
8392
}
@@ -90,12 +99,11 @@ Status ModelBuilder::RegisterInitializers() {
9099
const auto& name = tensor.name();
91100
const auto& shape = tensor.dims();
92101

93-
// Ignore the following tensors:
94-
// 1. Empty tensors: optional tensors can be indicated by an empty name.
95-
// 2. Tensors in skipped_initializers_: These are tensors that are not used as WebNN Constants.
96-
// Note: Scalar tensors are excluded because ONNX Runtime will optimize same scalar initializers into one.
97-
if (name.empty() || (Contains(skipped_initializers_, name) && !shape.empty()))
102+
// skip initializer if there is no remaining usage
103+
auto usage_count = initializer_usage_[name];
104+
if (usage_count == 0) {
98105
continue;
106+
}
99107

100108
std::vector<int32_t> dims;
101109
// 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
385393
}
386394

387395
void ModelBuilder::AddInitializerToSkip(const std::string& tensor_name) {
388-
skipped_initializers_.insert(tensor_name);
396+
// Decrement usage count if this is a known initializer.
397+
// For simplicity the OpBuilder::AddInitializersToSkip implementations may call this for arbitrary input names
398+
// without first checking if the value is an initializer.
399+
auto entry = initializer_usage_.find(tensor_name);
400+
if (entry != initializer_usage_.end()) {
401+
--entry->second;
402+
}
389403
}
390404

391405
void ModelBuilder::AddInputToSkip(const std::string& input_name) {

onnxruntime/core/providers/webnn/builders/model_builder.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ class ModelBuilder {
8181

8282
InlinedHashMap<std::string, OnnxTensorInfo> input_output_info_;
8383

84-
InlinedHashSet<std::string> skipped_initializers_;
84+
std::unordered_map<std::string, int> initializer_usage_;
8585
InlinedHashSet<std::string> skipped_inputs_;
8686

8787
uint32_t name_token_{0};

0 commit comments

Comments
 (0)