Add C++ IR (Graph) entry point to the optimizer

onnxoptimizer/optimize.cc

@@ -41,6 +41,14 @@ ModelProto OptimizeFixed(
   Optimizer current_opt(names, true);
   return current_opt.optimize(mp_in);
 }
+void OptimizeGraph(Graph& graph, const std::vector<std::string>& names) {
+  Optimizer current_opt(names, false);
+  current_opt.optimize(graph);
+}
+void OptimizeGraphFixed(Graph& graph, const std::vector<std::string>& names) {
+  Optimizer current_opt(names, true);
+  current_opt.optimize(graph);
+}
 const std::vector<std::string> GetAvailablePasses() {
   return Optimizer::passes.GetAvailablePasses();
 }

onnxoptimizer/optimize.h

@@ -26,6 +26,16 @@ struct Optimizer {
   Optimizer(const std::vector<std::string> &names, const bool fixed_point);
   ~Optimizer();
 
+  // Optimize the ONNX C++ IR (Graph) in place, running the configured passes
+  // directly on the graph. This avoids the ModelProto <-> Graph round-trip and
+  // is intended for C++ callers that already hold a Graph (e.g. one built by
+  // hand or produced by another IR pass). Proto-level concerns such as the
+  // ir_version upgrade and function copying are the caller's responsibility,
+  // since those live on ModelProto rather than on Graph.
+  void optimize(Graph &graph) {
+    this->pass_manager->run(graph);
+  }
+
   ModelProto optimize(const ModelProto &_mp_in) {
     const ModelProto* mp_in = &_mp_in;
     std::unique_ptr<ModelProto> copy_in;
@@ -46,7 +56,7 @@ struct Optimizer {
     }
 
     ModelProto mp_out = PrepareOutput(*mp_in);
-    this->pass_manager->run(*g);
+    this->optimize(*g);
     ExportModelProto(&mp_out, g);
 
     // Maybe we can optimize these functions, now just copy
@@ -99,5 +109,12 @@ ModelProto Optimize(const ModelProto &mp_in,
 
 ModelProto OptimizeFixed(const ModelProto &mp_in,
                          const std::vector<std::string> &names);
+
+// In-place counterparts that operate directly on the ONNX C++ IR (Graph),
+// skipping the ModelProto <-> Graph conversion. For C++ callers that already
+// hold a Graph.
+void OptimizeGraph(Graph &graph, const std::vector<std::string> &names);
+
+void OptimizeGraphFixed(Graph &graph, const std::vector<std::string> &names);
 }  // namespace optimization
 }  // namespace ONNX_NAMESPACE

tests/test_simple.cc

@@ -27,3 +27,39 @@ TEST(OptimizerTest, NopReshape) {
     ASSERT_EQ(optimized_model.graph().node().size(), 1);
     ASSERT_EQ(optimized_model.graph().node()[0].op_type(), "Identity");
 }
+
+// Exercises the in-place C++ IR (Graph) entry point: import a ModelProto into
+// the IR, optimize the Graph directly (no ModelProto round-trip inside the
+// optimizer), then export to confirm the passes mutated the graph in place.
+TEST(OptimizerTest, OptimizeGraphInPlace) {
+    const char* graph_str = R"(
+        <
+            ir_version: 7,
+            opset_import: [ "": 10]
+        >
+        agraph (float[5, 7] X) => (float[5, 7] Z)
+        {
+            Shape = Constant<value=int64[2]{5, -1}> ()
+            Y = Reshape (X, Shape)
+            Z = Identity(Y)
+        }
+    )";
+    onnx::ModelProto model;
+    const auto status = onnx::OnnxParser::Parse(model, graph_str);
+    EXPECT_TRUE(status.IsOK());
+
+    // ModelProto -> C++ IR.
+    std::shared_ptr<onnx::Graph> g(onnx::ImportModelProto(model));
+    ASSERT_NE(g.get(), nullptr);
+
+    // Optimize the IR in place.
+    onnx::optimization::OptimizeGraph(
+        *g, {"eliminate_nop_reshape", "eliminate_deadend"});
+
+    // Export the mutated IR to confirm the passes ran on the graph directly.
+    onnx::ModelProto optimized_model = onnx::PrepareOutput(model);
+    onnx::ExportModelProto(&optimized_model, g);
+
+    ASSERT_EQ(optimized_model.graph().node().size(), 1);
+    ASSERT_EQ(optimized_model.graph().node()[0].op_type(), "Identity");
+}