diff --git a/onnxruntime/core/providers/webnn/builders/helper.cc b/onnxruntime/core/providers/webnn/builders/helper.cc index d34cb7e362..7718fbdc2d 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.cc +++ b/onnxruntime/core/providers/webnn/builders/helper.cc @@ -12,6 +12,24 @@ namespace onnxruntime { namespace webnn { +InitializedTensorSet CollectAllInitializedTensors(const GraphViewer& graph_viewer) { + InitializedTensorSet all_initializers; + if (graph_viewer.IsSubgraph()) { + const Graph* cur_graph = &graph_viewer.GetGraph(); + // Traverse up to the top-level graph, collecting all initializers. + while (cur_graph->IsSubgraph()) { + const auto& current_initializers = cur_graph->GetAllInitializedTensors(); + all_initializers.insert(current_initializers.begin(), current_initializers.end()); + cur_graph = cur_graph->ParentGraph(); + } + // Collect initializers in top-level graph. + const auto& current_initializers = cur_graph->GetAllInitializedTensors(); + all_initializers.insert(current_initializers.begin(), current_initializers.end()); + } + + return all_initializers; +} + bool GetShape(const NodeArg& node_arg, std::vector& shape, const logging::Logger& logger) { const auto* shape_proto = node_arg.Shape(); if (!shape_proto) { diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 5aec81af15..ea57ab1af1 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -35,6 +35,9 @@ typedef struct { bool isCpuSupported; // The WebNN CPU backend XNNPack supports it (not about the CPU EP). } WebnnOpInfo; +// Collects all the initializer tensors in the subGraph and its ancestor graphs. +InitializedTensorSet CollectAllInitializedTensors(const GraphViewer& graph_viewer); + bool GetShape(const NodeArg& node_arg, std::vector& shape, const logging::Logger& logger); template diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.cc b/onnxruntime/core/providers/webnn/builders/model_builder.cc index b6631263df..b57e1b89b0 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/model_builder.cc @@ -38,6 +38,25 @@ Status ModelBuilder::Initialize() { return Status::OK(); } +InitializedTensorSet ModelBuilder::GetInitializerTensors() { + if (graph_viewer_.IsSubgraph()) { + auto all_initializers = CollectAllInitializedTensors(graph_viewer_); + const auto sub_graph_id = graph_viewer_.GetFilterInfo(); + const auto subgraph_initializer_names = sub_graph_id->GetMetaDef()->constant_initializers; + InitializedTensorSet subgraph_initializers; + + for (const auto& name : subgraph_initializer_names) { + auto it = all_initializers.find(name); + if (it != all_initializers.end()) { + subgraph_initializers.insert(*it); + } + } + return subgraph_initializers; + } else { + return graph_viewer_.GetAllInitializedTensors(); + } +} + /* static */ const IOpBuilder* ModelBuilder::GetOpBuilder(const Node& node) { const auto& op_builders = GetOpBuilders(); const auto it = op_builders.find(node.OpType()); diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.h b/onnxruntime/core/providers/webnn/builders/model_builder.h index c381eef3f4..16c8bf2d3c 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.h +++ b/onnxruntime/core/providers/webnn/builders/model_builder.h @@ -30,7 +30,7 @@ class ModelBuilder { // Accessors for members. const GraphViewer& GetGraphViewer() const { return graph_viewer_; } - const InitializedTensorSet& GetInitializerTensors() const { return graph_viewer_.GetAllInitializedTensors(); } + InitializedTensorSet GetInitializerTensors(); const emscripten::val& GetBuilder() const { return wnn_builder_; } const emscripten::val& GetContext() const { return wnn_context_; } diff --git a/onnxruntime/core/providers/webnn/webnn_execution_provider.cc b/onnxruntime/core/providers/webnn/webnn_execution_provider.cc index 4da54aaad3..cf18b3225e 100644 --- a/onnxruntime/core/providers/webnn/webnn_execution_provider.cc +++ b/onnxruntime/core/providers/webnn/webnn_execution_provider.cc @@ -59,10 +59,15 @@ WebNNExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_view const IKernelLookup& /*kernel_registries*/) const { std::vector> result; - // We do not run WebNN EP on subgraph, instead we cover this in the control flow nodes. - // TODO investigate whether we want to support subgraph using WebNN EP. - if (graph_viewer.IsSubgraph()) { - return result; + // For subgraph which is the attribute of the control flow nodes, part of its initializers are stored in its + // ancestor graphs as common initializers shared for other subgraphs. We need to collect all of them used for + // identifying the required initializer names and storing into 'meta_def->constant_initializers'. + // Thus we are able to get the required initialized tensors for this subgraph via the GetInitializerTensors() + // method defined in the model_builder.h file. + InitializedTensorSet all_initializers; + const bool is_subgraph = graph_viewer.IsSubgraph(); + if (is_subgraph) { + all_initializers = webnn::CollectAllInitializedTensors(graph_viewer); } /* @@ -110,6 +115,7 @@ WebNNExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_view std::unique_ptr sub_graph = std::make_unique(); + std::vector subgraph_initializers; InlinedHashSet node_outputs; InlinedHashSet subgraph_inputs; InlinedHashSet subgraph_outputs; @@ -126,7 +132,11 @@ WebNNExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_view // skip the placeholder inputs. continue; } - // if the node input was not produced by this subgraph, add it to the subgraph inputs. + // If it is a subgraph of a control flow node, collect the constant initializer. + if (is_subgraph && Contains(all_initializers, input->Name())) { + subgraph_initializers.push_back(input->Name()); + } + // If the node input was not produced by this subgraph, add it to the subgraph inputs. if (node_outputs.count(input) == 0) { if (subgraph_inputs.count(input) == 0) { subgraph_inputs.insert(input); @@ -165,6 +175,12 @@ WebNNExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_view meta_def->since_version = 1; meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL; + if (is_subgraph) { + for (const auto& initializer : subgraph_initializers) { + meta_def->constant_initializers.push_back(initializer); + } + } + for (const auto& input : ordered_subgraph_inputs) { meta_def->inputs.push_back(input->Name()); }