diff --git a/onnxoptimizer/optimize.cc b/onnxoptimizer/optimize.cc index f784bf608..cb4a7d792 100644 --- a/onnxoptimizer/optimize.cc +++ b/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& names) { + Optimizer current_opt(names, false); + current_opt.optimize(graph); +} +void OptimizeGraphFixed(Graph& graph, const std::vector& names) { + Optimizer current_opt(names, true); + current_opt.optimize(graph); +} const std::vector GetAvailablePasses() { return Optimizer::passes.GetAvailablePasses(); } diff --git a/onnxoptimizer/optimize.h b/onnxoptimizer/optimize.h index 6d59429e0..1fa086803 100644 --- a/onnxoptimizer/optimize.h +++ b/onnxoptimizer/optimize.h @@ -26,6 +26,16 @@ struct Optimizer { Optimizer(const std::vector &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 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 &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 &names); + +void OptimizeGraphFixed(Graph &graph, const std::vector &names); } // namespace optimization } // namespace ONNX_NAMESPACE diff --git a/tests/test_simple.cc b/tests/test_simple.cc index 64c1b3748..8ac22b508 100644 --- a/tests/test_simple.cc +++ b/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 () + 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 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"); +}