diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc index d0ae32378379d..12416ea0c121b 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.cc @@ -56,7 +56,13 @@ DEFINE_ADD_OPERAND_FROM_SCALAR(float, FLOAT32); #undef DEFINE_ADD_OPERAND_FROM_SCALAR 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 -= 1; + } } Status ModelBuilder::Prepare() { @@ -87,7 +93,16 @@ static size_t GetPaddedByteSize(size_t size) { } void ModelBuilder::PreprocessInitializers() { + const auto& initializers = GetInitializerTensors(); + for (const auto& node_unit : node_unit_holder_) { + // find all initializers consumed. AddInitializersToSkip will potentially decrement the usage count. + for (const auto& input : node_unit->Inputs()) { + if (input.node_arg.Exists() && Contains(initializers, input.node_arg.Name())) { + initializer_usage_[input.node_arg.Name()]++; + } + } + if (const auto* op_builder = GetOpBuilder(*node_unit)) { op_builder->AddInitializersToSkip(*this, *node_unit); } @@ -208,11 +223,16 @@ Status ModelBuilder::RegisterInitializers() { std::vector> initializers(initializer_size); size_t sizeAll = 0; + const auto should_skip_initializer = [this](const std::string& name) -> bool { + const auto it = initializer_usage_.find(name); + return it == initializer_usage_.end() || it->second == 0; + }; + int i = 0; for (const auto& pair : initializer_tensors) { const auto& tensor = *pair.second; const auto& name = tensor.name(); - if (Contains(skipped_initializers_, name)) + if (should_skip_initializer(name)) continue; Shape shape; @@ -249,7 +269,7 @@ Status ModelBuilder::RegisterInitializers() { size_t offset = 0; for (const auto& pair : initializer_tensors) { const auto& tensor = *pair.second; - if (Contains(skipped_initializers_, tensor.name())) + if (should_skip_initializer(tensor.name())) continue; auto [index, size, padded_size] = initializers[i++]; @@ -439,10 +459,11 @@ Status ModelBuilder::AddOperandFromPersistMemoryBuffer( Status ModelBuilder::AddOperations() { const auto& node_indices = graph_viewer_.GetNodesInTopologicalOrder(); for (const auto node_idx : node_indices) { - LOGS_DEFAULT(VERBOSE) << "Adding node [" << node_idx << "]"; const auto* node(graph_viewer_.GetNode(node_idx)); const NodeUnit& node_unit = GetNodeUnit(node); + LOGS_DEFAULT(VERBOSE) << "Adding node [" << node_unit.Name() << "] at index [" << node_unit.Index() << "]"; + // Since we may have NodeUnit with multiple nodes, insert NodeUnit with the first occurrence of // its node(s) in topological order may cause the incorrect topological order while inserting // NodeUNits, for example, diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.h index 8eddf389d3b52..b2118150dd304 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/model_builder.h @@ -134,7 +134,7 @@ class ModelBuilder { std::unordered_set operands_; std::unordered_set fused_activations_; - std::unordered_set skipped_initializers_; + std::unordered_map initializer_usage_; // All activation nodes (Relu, Relu1, Relu6) as a map std::unordered_map activation_node_units_;