Add optional graph-modification report to optimize()

onnxoptimizer/__init__.py

@@ -33,16 +33,27 @@
 
 
 def optimize(
-    model: onnx.ModelProto, passes: list[str] | None = None, fixed_point: bool = False
-) -> onnx.ModelProto:
+    model: onnx.ModelProto,
+    passes: list[str] | None = None,
+    fixed_point: bool = False,
+    return_report: bool = False,
+) -> onnx.ModelProto | tuple[onnx.ModelProto, dict[str, int]]:
     """Apply the optimization on the serialized ModelProto.
 
     Arguments:
         model: ONNX model.
         passes: Optimization names.
+        fixed_point: Whether to repeatedly apply the passes until the graph
+            reaches a fixed point.
+        return_report: When ``True``, also return a report describing how many
+            times each pass modified the graph.
 
     Return:
-        Optimized model.
+        The optimized model. When ``return_report`` is ``True``, a tuple of
+        ``(optimized_model, report)`` is returned instead, where ``report`` is a
+        ``dict`` mapping each pass name to the total number of positive
+        (successful) transforms it applied. A pass that ran without changing the
+        graph maps to ``0``; passes absent from the map did not report a count.
     """
 
     if passes is None:
@@ -51,6 +62,14 @@ def optimize(
         raise TypeError(f"Optimizer only accepts ModelProto, incorrect type: {type(model)}")
     try:
         model_str = model.SerializeToString()
+        if return_report:
+            if fixed_point:
+                optimized_model_str, report = _c.optimize_fixedpoint_report(
+                    model_str, passes
+                )
+            else:
+                optimized_model_str, report = _c.optimize_report(model_str, passes)
+            return onnx.load_from_string(optimized_model_str), report
         if fixed_point:
             optimized_model_str = _c.optimize_fixedpoint(model_str, passes)
         else:
@@ -78,15 +97,28 @@ def optimize(
                     location=data_src_rel_filename,
                     convert_attribute=True,
                 )
-                if fixed_point:
+                report = None
+                if return_report:
+                    if fixed_point:
+                        report = _c.optimize_fixedpoint_from_path_report(
+                            file_src.name, file_dest.name, passes, data_dest_rel_filename
+                        )
+                    else:
+                        report = _c.optimize_from_path_report(
+                            file_src.name, file_dest.name, passes, data_dest_rel_filename
+                        )
+                elif fixed_point:
                     _c.optimize_fixedpoint_from_path(
                         file_src.name, file_dest.name, passes, data_dest_rel_filename
                     )
                 else:
                     _c.optimize_from_path(
                         file_src.name, file_dest.name, passes, data_dest_rel_filename
                     )
-                return onnx.load(file_dest.name, load_external_data=True)
+                optimized_model = onnx.load(file_dest.name, load_external_data=True)
+                if return_report:
+                    return optimized_model, report
+                return optimized_model
             finally:
                 os.remove(file_src.name)
                 os.remove(file_dest.name)

onnxoptimizer/cpp2py_export.cc

@@ -3,9 +3,13 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/map.h>
+#include <nanobind/stl/pair.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/vector.h>
 
+#include <map>
+
 #include "onnxoptimizer/model_util.h"
 #include "onnxoptimizer/optimize.h"
 
@@ -48,6 +52,30 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
         return nb::bytes(out.data(), out.size());
       });
 
+  onnx_opt_cpp2py_export.def(
+      "optimize_report",
+      [](const nb::bytes& bytes, const std::vector<std::string>& names) {
+        ModelProto proto{};
+        ParseProtoFromPyBytes(&proto, bytes);
+        std::map<std::string, unsigned int> report;
+        auto const result = optimization::Optimize(proto, names, &report);
+        std::string out;
+        result.SerializeToString(&out);
+        return std::make_pair(nb::bytes(out.data(), out.size()), report);
+      });
+
+  onnx_opt_cpp2py_export.def(
+      "optimize_fixedpoint_report",
+      [](const nb::bytes& bytes, const std::vector<std::string>& names) {
+        ModelProto proto{};
+        ParseProtoFromPyBytes(&proto, bytes);
+        std::map<std::string, unsigned int> report;
+        auto const result = optimization::OptimizeFixed(proto, names, &report);
+        std::string out;
+        result.SerializeToString(&out);
+        return std::make_pair(nb::bytes(out.data(), out.size()), report);
+      });
+
   onnx_opt_cpp2py_export.def(
       "optimize_from_path", [](const std::string& import_model_path,
                                const std::string& export_model_path,
@@ -60,6 +88,21 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
                                 export_data_file_name);
       });
 
+  onnx_opt_cpp2py_export.def(
+      "optimize_from_path_report",
+      [](const std::string& import_model_path,
+         const std::string& export_model_path,
+         const std::vector<std::string>& names,
+         const std::string& export_data_file_name) {
+        ModelProto proto{};
+        optimization::loadModel(&proto, import_model_path, true);
+        std::map<std::string, unsigned int> report;
+        auto result = optimization::Optimize(proto, names, &report);
+        optimization::saveModel(&result, export_model_path, true,
+                                export_data_file_name);
+        return report;
+      });
+
   onnx_opt_cpp2py_export.def(
       "optimize_fixedpoint_from_path",
       [](const std::string& import_model_path,
@@ -72,6 +115,21 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
         optimization::saveModel(&result, export_model_path, true,
                                 export_data_file_name);
       });
+
+  onnx_opt_cpp2py_export.def(
+      "optimize_fixedpoint_from_path_report",
+      [](const std::string& import_model_path,
+         const std::string& export_model_path,
+         const std::vector<std::string>& names,
+         const std::string& export_data_file_name) {
+        ModelProto proto{};
+        optimization::loadModel(&proto, import_model_path, true);
+        std::map<std::string, unsigned int> report;
+        auto result = optimization::OptimizeFixed(proto, names, &report);
+        optimization::saveModel(&result, export_model_path, true,
+                                export_data_file_name);
+        return report;
+      });
   onnx_opt_cpp2py_export.def("get_available_passes",
                              &optimization::GetAvailablePasses);
   onnx_opt_cpp2py_export.def("get_fuse_and_elimination_passes",

onnxoptimizer/optimize.cc

@@ -31,15 +31,17 @@ Optimizer::~Optimizer() {}
 
 ModelProto Optimize(
     const ModelProto& mp_in,
-    const std::vector<std::string>& names) {
+    const std::vector<std::string>& names,
+    std::map<std::string, unsigned int>* report) {
   Optimizer current_opt(names, false);
-  return current_opt.optimize(mp_in);
+  return current_opt.optimize(mp_in, report);
 }
 ModelProto OptimizeFixed(
     const ModelProto& mp_in,
-    const std::vector<std::string>& names) {
+    const std::vector<std::string>& names,
+    std::map<std::string, unsigned int>* report) {
   Optimizer current_opt(names, true);
-  return current_opt.optimize(mp_in);
+  return current_opt.optimize(mp_in, report);
 }
 const std::vector<std::string> GetAvailablePasses() {
   return Optimizer::passes.GetAvailablePasses();

onnxoptimizer/optimize.h

@@ -26,7 +26,10 @@ struct Optimizer {
   Optimizer(const std::vector<std::string> &names, const bool fixed_point);
   ~Optimizer();
 
-  ModelProto optimize(const ModelProto &_mp_in) {
+  // If `report` is non-null it is filled with a map from pass name to the
+  // total number of positive transforms that pass applied to the graph.
+  ModelProto optimize(const ModelProto &_mp_in,
+                      std::map<std::string, unsigned int> *report = nullptr) {
     const ModelProto* mp_in = &_mp_in;
     std::unique_ptr<ModelProto> copy_in;
     if (mp_in->ir_version() == 3) {
@@ -46,7 +49,10 @@ struct Optimizer {
     }
 
     ModelProto mp_out = PrepareOutput(*mp_in);
-    this->pass_manager->run(*g);
+    auto analysis = this->pass_manager->run(*g);
+    if (report != nullptr && analysis != nullptr) {
+      *report = analysis->transform_counts;
+    }
     ExportModelProto(&mp_out, g);
 
     // Maybe we can optimize these functions, now just copy
@@ -95,9 +101,11 @@ const std::vector<std::string> GetAvailablePasses();
 const std::vector<std::string> GetFuseAndEliminationPass();
 
 ModelProto Optimize(const ModelProto &mp_in,
-                    const std::vector<std::string> &names);
+                    const std::vector<std::string> &names,
+                    std::map<std::string, unsigned int> *report = nullptr);
 
 ModelProto OptimizeFixed(const ModelProto &mp_in,
-                         const std::vector<std::string> &names);
+                         const std::vector<std::string> &names,
+                         std::map<std::string, unsigned int> *report = nullptr);
 }  // namespace optimization
 }  // namespace ONNX_NAMESPACE

onnxoptimizer/pass_manager.cc

@@ -19,14 +19,23 @@ void GeneralPassManager::add(std::shared_ptr<Pass> pass) {
 }
 
 std::shared_ptr<PassManagerAnalysis> GeneralPassManager::run(Graph& graph) {
+  auto report = std::make_shared<PassManagerAnalysis>();
   for (const std::shared_ptr<Pass>& pass : this->passes) {
-    auto pass_analysis = pass->runPass(graph);
+    std::shared_ptr<PostPassAnalysis> analysis = pass->runPass(graph);
+    if (pass->getPassAnalysisType() == PassAnalysisType::Empty) {
+      continue;
+    }
+    std::shared_ptr<CountBasedPassAnalysis> count_analysis =
+        std::static_pointer_cast<CountBasedPassAnalysis>(analysis);
+    report->transform_counts[pass->getPassName()] +=
+        count_analysis->num_positive_transforms;
   }
-  return std::shared_ptr<PassManagerAnalysis>(new EmptyPassManagerAnalysis());
+  return report;
 }
 
 std::shared_ptr<PassManagerAnalysis> FixedPointPassManager::run(Graph& graph) {
   bool fixed_point_optimization_done;
+  auto report = std::make_shared<PassManagerAnalysis>();
 
   do {
     fixed_point_optimization_done = false;
@@ -37,13 +46,17 @@ std::shared_ptr<PassManagerAnalysis> FixedPointPassManager::run(Graph& graph) {
       }
       std::shared_ptr<CountBasedPassAnalysis> count_analysis =
           std::static_pointer_cast<CountBasedPassAnalysis>(analysis);
+      report->transform_counts[pass->getPassName()] +=
+          count_analysis->num_positive_transforms;
       if (count_analysis->num_positive_transforms != 0) {
         VLOG(1) << "Pass " << pass->getPassName() << " transformed " << count_analysis->num_positive_transforms;
       }
 
       while (count_analysis->fixedPointOptimizationNeeded()) {
         count_analysis = std::static_pointer_cast<CountBasedPassAnalysis>(
             pass->runPass(graph));
+        report->transform_counts[pass->getPassName()] +=
+            count_analysis->num_positive_transforms;
         if (count_analysis->num_positive_transforms != 0) {
           VLOG(1) << "Pass " << pass->getPassName() << " transformed " << count_analysis->num_positive_transforms;
         }
@@ -52,7 +65,7 @@ std::shared_ptr<PassManagerAnalysis> FixedPointPassManager::run(Graph& graph) {
     }
   } while (fixed_point_optimization_done);
 
-  return std::shared_ptr<PassManagerAnalysis>(new EmptyPassManagerAnalysis());
+  return report;
 }
 } // namespace optimization
 } // namespace ONNX_NAMESPACE

onnxoptimizer/pass_manager.h

@@ -6,14 +6,23 @@
 // ATTENTION: The code in this file is highly EXPERIMENTAL.
 // Adventurous users should note that the APIs will probably change.
 
+#include <map>
+#include <string>
 #include <vector>
 #include "onnxoptimizer/pass.h"
 
 namespace ONNX_NAMESPACE {
 namespace optimization {
 
 // An analysis returned from the run done by a manager
-struct PassManagerAnalysis {};
+struct PassManagerAnalysis {
+  virtual ~PassManagerAnalysis() = default;
+  // Maps a pass name to the total number of positive (successful) transforms
+  // that pass applied to the graph during the run. A pass is only present in
+  // the map when it was given a chance to report a count (i.e. count-based
+  // passes); a value of 0 means the pass ran but did not modify the graph.
+  std::map<std::string, unsigned int> transform_counts;
+};
 struct EmptyPassManagerAnalysis : PassManagerAnalysis {};
 
 // Base class of all PassManager's. The class should be able to add new passes

onnxoptimizer/test/optimizer_test.py

@@ -370,6 +370,66 @@ def check_identity(node):  # type: (NodeProto) -> None
         assert len(optimized_model.graph.node[3].attribute[0].g.output) == 2
         assert optimized_model.graph.node[3].attribute[0].g.output[1].name == "_B2"
 
+    def test_optimize_return_report(self):  # type: () -> None
+        # `return_report=True` should return a (model, report) tuple, where the
+        # report maps each pass name to how many times it modified the graph.
+        nodes = [
+            helper.make_node("Add", ["X", "Y"], ["A"]),
+            helper.make_node("Identity", ["A"], ["B"]),
+        ]
+        graph = helper.make_graph(
+            nodes,
+            "test",
+            [
+                helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,)),
+                helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5,)),
+            ],
+            [helper.make_tensor_value_info("B", TensorProto.FLOAT, (5,))],
+        )
+        model = helper.make_model(
+            graph,
+            producer_name="onnx-test",
+            opset_imports=[helper.make_opsetid("", LATEST_STABLE_OPSET_VERSION)],
+            ir_version=10,
+        )
+
+        result = onnxoptimizer.optimize(
+            model, ["eliminate_identity"], return_report=True
+        )
+        assert isinstance(result, tuple)
+        optimized_model, report = result
+        assert isinstance(optimized_model, ModelProto)
+        assert isinstance(report, dict)
+        # The single Identity node should have been eliminated exactly once.
+        assert report.get("eliminate_identity", 0) == 1
+
+        # Without return_report the return value stays a bare model (backward compat).
+        optimized_only = onnxoptimizer.optimize(model, ["eliminate_identity"])
+        assert isinstance(optimized_only, ModelProto)
+
+    def test_optimize_return_report_no_change(self):  # type: () -> None
+        # A pass that matches nothing should still be reportable and count 0.
+        node = helper.make_node("Add", ["X", "Y"], ["Z"])
+        graph = helper.make_graph(
+            [node],
+            "test",
+            [
+                helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,)),
+                helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5,)),
+            ],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5,))],
+        )
+        model = helper.make_model(
+            graph,
+            producer_name="onnx-test",
+            opset_imports=[helper.make_opsetid("", LATEST_STABLE_OPSET_VERSION)],
+            ir_version=10,
+        )
+        _, report = onnxoptimizer.optimize(
+            model, ["eliminate_identity"], return_report=True
+        )
+        assert report.get("eliminate_identity", 0) == 0
+
     # type: () -> None
     def test_eliminate_identity_both_graph_input_and_output(self):
         # We should not eliminate an op when its input is also graph input,