diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index a4e64bc476..d1a3f6bd54 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -966,7 +966,6 @@ class Graph { ONNX_NAMESPACE::GraphProto* graph_proto_; InitializedTensorSet name_to_initial_tensor_; - std::vector removed_initializer_indexes_; IOnnxRuntimeOpSchemaCollectionPtr schema_registry_; diff --git a/onnxruntime/core/graph/function.cc b/onnxruntime/core/graph/function.cc index 785bc6d938..0fed268f89 100644 --- a/onnxruntime/core/graph/function.cc +++ b/onnxruntime/core/graph/function.cc @@ -207,7 +207,11 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, for (const auto& input : meta_def->inputs) { const ONNX_NAMESPACE::TensorProto* initializer = nullptr; if (graph.GetInitializedTensor(input, initializer)) { - function_body_graph.AddInitializedTensor(*initializer); + // meta_def->inputs could have duplicates so make sure we only add once + const ONNX_NAMESPACE::TensorProto* subgraph_initializer = nullptr; + if (!function_body_graph.GetInitializedTensor(input, subgraph_initializer)) { + function_body_graph.AddInitializedTensor(*initializer); + } } } @@ -221,9 +225,9 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, const ONNX_NAMESPACE::FunctionProto& onnx_func_proto) : parent_graph_(&graph) { // Make a copy of the FunctionProto. - // All FunctionBody ops with the same op type seem to share the same FunctionProto struct within a model. + // All FunctionBody ops with the same op type seem to share the same FunctionProto struct within a model. // Hence, we make a copy prior to generating the graph representation of the function, - // as we might make some modifications to the FunctionProto along the way + // as we might make some modifications to the FunctionProto along the way onnx_func_proto_ = onnx_func_proto; auto node_in_parent_graph = parent_graph_->GetNode(node_index); @@ -360,12 +364,13 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph, new_attr_map[(*node_attr).name()] = *node_attr; } } - function_body_graph.AddNode(uniq_identifier + "_" + std::to_string(node_index), (*node).op_type(), (*node).doc_string(), inputs, outputs, &new_attr_map, (*node).domain()); + function_body_graph.AddNode(uniq_identifier + "_" + std::to_string(node_index), (*node).op_type(), + (*node).doc_string(), inputs, outputs, &new_attr_map, (*node).domain()); } auto status = function_body_graph.Resolve(); ORT_ENFORCE(status.IsOK(), "Resolve subgraph failed:", status.ErrorMessage()); -} // namespace onnxruntime +} FunctionImpl::~FunctionImpl() = default; diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 9c02765352..b31cc915be 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2029,7 +2029,10 @@ void Graph::SetDescription(const std::string& description) { } void Graph::AddInitializedTensor(const TensorProto& tensor) { - if (name_to_initial_tensor_.end() != name_to_initial_tensor_.find(tensor.name())) { + auto existing = name_to_initial_tensor_.find(tensor.name()); + if (existing != name_to_initial_tensor_.cend()) { + ORT_ENFORCE(existing->second == &tensor, + "AddInitializedTensor already has tensor with name ", tensor.name(), " but different TensorProto."); return; } @@ -2046,25 +2049,43 @@ void Graph::AddInitializedTensor(const TensorProto& tensor) { ORT_IGNORE_RETURN_VALUE(GetOrCreateNodeArg(tensor.name(), &t)); } - - SetGraphProtoSyncNeeded(); - SetGraphResolveNeeded(); } void Graph::RemoveInitializedTensor(const std::string& tensor_name) { + bool found = false; auto iter = name_to_initial_tensor_.find(tensor_name); - if (name_to_initial_tensor_.end() != iter) { + found = iter != name_to_initial_tensor_.end(); + if (found) { name_to_initial_tensor_.erase(tensor_name); - SetGraphProtoSyncNeeded(); SetGraphResolveNeeded(); } + + auto& mutable_initializers = *(graph_proto_->mutable_initializer()); + auto proto_entry = std::find_if(mutable_initializers.begin(), mutable_initializers.end(), + [&tensor_name](const TensorProto& entry) { return entry.name() == tensor_name; }); + + if (proto_entry != mutable_initializers.end()) { + auto num_entries = mutable_initializers.size(); + if (num_entries > 1) { + // swap the entry being deleted with the last one, and delete it. + // we do this so we don't have to move all the entries past the one being deleted down one. + auto slot = proto_entry - mutable_initializers.begin(); + auto last_entry = mutable_initializers.end() - 1; + mutable_initializers.SwapElements(slot, num_entries - 1); + mutable_initializers.erase(last_entry); + } else { + mutable_initializers.erase(proto_entry); + } + } else { + // these should always be in sync as the pointer in name_to_initial_tensor_ is to memory owned by graph_proto_ + ORT_ENFORCE(!found, "graph_proto_ is not in sync with name_to_initial_tensor_."); + } } Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_initializer) { // name_to_initial_tensor_ maps from name to const TensorProto*, so we first // look up the const pointer by name, then find and modify the mutable // pointed-to TensorProto in graph_proto_. - const auto& initializer_name = new_initializer.name(); const auto name_to_initializer_it = name_to_initial_tensor_.find(initializer_name); ORT_RETURN_IF_NOT(name_to_initializer_it != name_to_initial_tensor_.end(), @@ -2085,13 +2106,15 @@ Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_in "Replacement tensor's data type does not match."); auto& mutable_initializers = *(graph_proto_->mutable_initializer()); - auto old_mutable_initializer_ptr_it = std::find( - mutable_initializers.pointer_begin(), mutable_initializers.pointer_end(), &old_initializer); - ORT_ENFORCE(old_mutable_initializer_ptr_it != mutable_initializers.pointer_end(), - "graph_proto_ is not in sync with name_to_initial_tensor_"); - auto& old_mutable_initializer = **old_mutable_initializer_ptr_it; + // use cheaper pointer comparison to find old entry + auto existing_entry = std::find(mutable_initializers.pointer_begin(), mutable_initializers.pointer_end(), + &old_initializer); - old_mutable_initializer = new_initializer; + // these should always be in sync as the pointer in name_to_initial_tensor_ is to memory owned by graph_proto_ + ORT_ENFORCE(existing_entry != mutable_initializers.pointer_end(), + "graph_proto_ is not in sync with name_to_initial_tensor_"); + + **existing_entry = new_initializer; return Status::OK(); } @@ -2108,7 +2131,6 @@ bool Graph::GetInitializedTensor(const std::string& tensor_name, const TensorPro void Graph::CleanAllInitializedTensors() noexcept { name_to_initial_tensor_.clear(); - removed_initializer_indexes_.clear(); // Clearing RepeatedPtrFields does not free objects' memory. The memory is retained // and can be reused. Need to explicitly release the cleared objects and free the @@ -2288,35 +2310,6 @@ const ONNX_NAMESPACE::GraphProto& Graph::ToGraphProto() { // Nodes. ToGraphProtoInternal(*graph_proto_); - if (!removed_initializer_indexes_.empty()) { - // Move initializers. - std::sort(removed_initializer_indexes_.begin(), removed_initializer_indexes_.end()); - int lastInUseInitializerIndex = graph_proto_->initializer_size() - 1; - int start = 0; - int end = gsl::narrow_cast(removed_initializer_indexes_.size()) - 1; - int lastRemovedInitializerIndex = removed_initializer_indexes_[end]; - - for (; start <= end; start++) { - // Find a lastInUseInitializer. - while (start <= end && lastInUseInitializerIndex == lastRemovedInitializerIndex) { - graph_proto_->mutable_initializer()->RemoveLast(); - lastInUseInitializerIndex--; - end--; - if (start <= end) { - lastRemovedInitializerIndex = removed_initializer_indexes_[end]; - } - } - - if (start <= end) { - // Copy the initializer in use to the slot which is removed. - *graph_proto_->mutable_initializer(removed_initializer_indexes_[start]) = graph_proto_->initializer(lastInUseInitializerIndex); - graph_proto_->mutable_initializer()->RemoveLast(); - lastInUseInitializerIndex--; - } - } - removed_initializer_indexes_.clear(); - } - GraphProtoSyncNeeded(false); return *graph_proto_; @@ -2329,9 +2322,7 @@ ONNX_NAMESPACE::GraphProto Graph::ToGraphProto() const { GraphProto result; ToGraphProtoInternal(result); - for (auto initializer : GetAllInitializedTensors()) { - *result.add_initializer() = *initializer.second; - } + *result.mutable_initializer() = graph_proto_->initializer(); return result; } diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index 06b86a0124..530bf37995 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -959,8 +959,7 @@ TEST(GraphUpdateTest, ReplaceInitializedTensor) { status = graph.ReplaceInitializedTensor(valid_replacement); ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); - auto tensor_data_matches = []( - const ONNX_NAMESPACE::TensorProto& a, const ONNX_NAMESPACE::TensorProto& b) { + auto tensor_data_matches = [](const ONNX_NAMESPACE::TensorProto& a, const ONNX_NAMESPACE::TensorProto& b) { if (a.int32_data_size() != b.int32_data_size()) return false; for (int i = 0; i < a.int32_data_size(); ++i) { if (a.int32_data(i) != b.int32_data(i)) return false; @@ -979,5 +978,59 @@ TEST(GraphUpdateTest, ReplaceInitializedTensor) { ASSERT_TRUE(tensor_data_matches(graph_proto.initializer(0), valid_replacement)); } } + +TEST(GraphUpdateTest, AddRemoveInitializerHandling) { + Model m{"test_model"}; + Graph& graph = m.MainGraph(); + + auto create_tensor_proto = [](const std::string& name, int32_t value) { + ONNX_NAMESPACE::TensorProto init{}; + init.set_name(name); + init.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + init.add_dims(1); + init.add_int32_data(value); + + return init; + }; + + auto init = create_tensor_proto("1", 1); + auto init2 = create_tensor_proto("2", 2); + + // add both, remove the 1st (moves the second initializer into the first slot), and finally re-add the first + graph.AddInitializedTensor(init); + graph.AddInitializedTensor(init2); + graph.RemoveInitializedTensor(init.name()); + graph.AddInitializedTensor(init); + + ASSERT_EQ(graph.GetAllInitializedTensors().size(), 2); + + // check the values coming from name_to_initial_tensor_ are good; + const TensorProto* i = nullptr; + ASSERT_TRUE(graph.GetInitializedTensor(init.name(), i)); + ASSERT_TRUE(i->int32_data()[0] == 1); + ASSERT_TRUE(graph.GetInitializedTensor(init2.name(), i)); + ASSERT_TRUE(i->int32_data()[0] == 2); + + // check the values in the GraphProto are also correct + ONNX_NAMESPACE::GraphProto graph_proto_from_const_graph = static_cast(graph).ToGraphProto(); + ONNX_NAMESPACE::GraphProto graph_proto_from_graph = graph.ToGraphProto(); + + ASSERT_EQ(graph_proto_from_const_graph.initializer_size(), 2); + ASSERT_EQ(graph_proto_from_graph.initializer_size(), 2); + + auto validate_proto = [&](const GraphProto& proto) { + auto initializers = proto.initializer(); + // we expect '2' to be before '1' due to the remove moving the last initializer into the slot of the one being + // removed in order to free memory and only move one entry + EXPECT_EQ(initializers[0].name(), init2.name()); + EXPECT_EQ(initializers[0].int32_data()[0], 2); + + EXPECT_EQ(initializers[1].name(), init.name()); + EXPECT_EQ(initializers[1].int32_data()[0], 1); + }; + + validate_proto(graph_proto_from_const_graph); + validate_proto(graph_proto_from_graph); +} } // namespace test } // namespace onnxruntime