add GetMutableMetaDef

microsoft · Feb 9, 2025 · df5aca9 · df5aca9
1 parent d4968cb
commit df5aca9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 17 deletions.
diff --git a/include/onnxruntime/core/graph/indexed_sub_graph.h b/include/onnxruntime/core/graph/indexed_sub_graph.h
@@ -70,6 +70,12 @@ struct IndexedSubGraph {
     return meta_def_.get();
   }
 
+  /** Gets the mutable meta definition needed to represent this subgraph as a FunctionProto.
+  @returns MetaDef instance if it has been set. nullptr if not. */
+  MetaDef* GetMutableMetaDef() const {
+    return meta_def_.get();
+  }
+
  private:
   // subgraph meta definition.
   std::unique_ptr<MetaDef> meta_def_;

diff --git a/onnxruntime/core/optimizer/selection_and_optimization_func.cc b/onnxruntime/core/optimizer/selection_and_optimization_func.cc
@@ -86,31 +86,21 @@ Status ConstantFoldingDQ_optimization(Graph& graph, const ComputeCapability& opt
   }
   cc_to_update.sub_graph->nodes = updated_nodes;
 
-  auto original_meta_def = cc_to_update.sub_graph->GetMetaDef();
-  std::unique_ptr<IndexedSubGraph::MetaDef> updated_meta_def = std::make_unique<IndexedSubGraph::MetaDef>();
-  updated_meta_def->name = original_meta_def->name;
-  updated_meta_def->domain = original_meta_def->domain;
-  updated_meta_def->since_version = original_meta_def->since_version;
-  updated_meta_def->status = original_meta_def->status;
-  updated_meta_def->inputs = original_meta_def->inputs;
-  updated_meta_def->outputs = original_meta_def->outputs;
-  updated_meta_def->attributes = original_meta_def->attributes;
-  updated_meta_def->doc_string = original_meta_def->doc_string;
-#if !defined(ORT_MINIMAL_BUILD)
-  updated_meta_def->type_and_shape_inference_function = original_meta_def->type_and_shape_inference_function;
-#endif
-  for (auto constant_initializer : original_meta_def->constant_initializers) {
+  auto meta_def = cc_to_update.sub_graph->GetMutableMetaDef();
+  std::vector<std::string> updated_constant_initializers;
+
+  for (auto constant_initializer : meta_def->constant_initializers) {
     if (original_initializers_to_remove.find(constant_initializer) != original_initializers_to_remove.end()) {
       continue;
     }
-    updated_meta_def->constant_initializers.push_back(constant_initializer);
+    updated_constant_initializers.push_back(constant_initializer);
   }
 
   for (auto constant_initializer : new_initializers_to_add) {
-    updated_meta_def->constant_initializers.push_back(constant_initializer);
+    updated_constant_initializers.push_back(constant_initializer);
   }
 
-  cc_to_update.sub_graph->SetMetaDef(std::move(updated_meta_def));
+  meta_def->constant_initializers = updated_constant_initializers;
 
   return Status::OK();
 }