eliminate_deadend/eliminate_common_subexpression: avoid uses()'s O(N) capture scan

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "third_party/onnx"]
 	path = third_party/onnx
-	url = https://github.com/onnx/onnx.git
+	url = https://github.com/onnxsim/onnx.git
+[submodule "third_party/blake3"]
+	path = third_party/blake3
+	url = https://github.com/BLAKE3-team/BLAKE3.git

CMakeLists.txt

@@ -48,8 +48,30 @@ file(GLOB onnx_opt_srcs CONFIGURE_DEPENDS "onnxoptimizer/*.cc"
     )
 list(REMOVE_ITEM onnx_opt_srcs "${PROJECT_SOURCE_DIR}/onnxoptimizer/cpp2py_export.cc")
 
+# BLAKE3 (third_party/blake3, a git submodule pinned to a specific upstream
+# commit -- see .gitmodules / `git submodule status`, and LICENSE_A2 /
+# LICENSE_CC0 in that directory) backing tensor_content_hash.cc's
+# TensorContentDigest, used by cse_util.h's CSETensorHash/CSETensorCompare.
+# Only the portable C sources (under the submodule's own c/ subdirectory --
+# upstream's layout) are compiled, with every SIMD backend disabled: that
+# makes blake3_dispatch.c's runtime CPU-feature dispatch fall through to
+# blake3_portable.c unconditionally, which is what makes omitting the SIMD
+# source files (and their per-architecture assembly) correct rather than
+# merely omitted. Mirrors onnxsim's own onnxsim/tensor_pool_hash.cpp
+# vendoring of the same library.
+set(ONNXOPT_BLAKE3_SOURCES
+    third_party/blake3/c/blake3.c
+    third_party/blake3/c/blake3_dispatch.c
+    third_party/blake3/c/blake3_portable.c)
+set(ONNXOPT_BLAKE3_COMPILE_DEFS
+    BLAKE3_NO_SSE2 BLAKE3_NO_SSE41 BLAKE3_NO_AVX2 BLAKE3_NO_AVX512
+    BLAKE3_USE_NEON=0)
+list(APPEND onnx_opt_srcs ${ONNXOPT_BLAKE3_SOURCES})
+
 onnxopt_add_library(onnx_optimizer ${onnx_opt_srcs})
 target_link_libraries(onnx_optimizer PUBLIC ${ONNX_TARGET_NAME})
+target_compile_definitions(onnx_optimizer PRIVATE ${ONNXOPT_BLAKE3_COMPILE_DEFS})
+target_include_directories(onnx_optimizer PRIVATE ${PROJECT_SOURCE_DIR}/third_party)
 target_include_directories(onnx_optimizer PUBLIC
     $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>
     $<INSTALL_INTERFACE:include>

onnxoptimizer/__init__.py

@@ -33,13 +33,23 @@
 
 
 def optimize(
-    model: onnx.ModelProto, passes: list[str] | None = None, fixed_point: bool = False
+    model: onnx.ModelProto,
+    passes: list[str] | None = None,
+    fixed_point: bool = False,
+    initializers_as_constants: bool = True,
 ) -> onnx.ModelProto:
     """Apply the optimization on the serialized ModelProto.
 
     Arguments:
         model: ONNX model.
         passes: Optimization names.
+        fixed_point: Whether to run the passes to a fixed point.
+        initializers_as_constants: Whether the passes may treat graph
+            initializers as constant tensors (the default, ``True``). When set
+            to ``False`` initializers are treated as non-constant, so
+            value-baking passes such as ``fuse_bn_into_conv`` leave
+            initializer-backed weights untouched; ``Constant`` nodes are still
+            treated as constants.
 
     Return:
         Optimized model.
@@ -49,6 +59,20 @@ def optimize(
         passes = get_fuse_and_elimination_passes()
     if not isinstance(model, onnx.ModelProto):
         raise TypeError(f"Optimizer only accepts ModelProto, incorrect type: {type(model)}")
+    # The C++ core reads this switch from thread-local state deep inside the
+    # passes, so set it around the call and restore it afterwards to avoid
+    # leaking the setting to unrelated callers on the same thread.
+    previous = _c.initializers_as_constants()
+    _c.set_initializers_as_constants(initializers_as_constants)
+    try:
+        return _optimize_impl(model, passes, fixed_point)
+    finally:
+        _c.set_initializers_as_constants(previous)
+
+
+def _optimize_impl(
+    model: onnx.ModelProto, passes: list[str], fixed_point: bool
+) -> onnx.ModelProto:
     try:
         model_str = model.SerializeToString()
         if fixed_point:

onnxoptimizer/cpp2py_export.cc

@@ -31,7 +31,13 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
       [](const nb::bytes& bytes, const std::vector<std::string>& names) {
         ModelProto proto{};
         ParseProtoFromPyBytes(&proto, bytes);
-        auto const result = optimization::Optimize(proto, names);
+        // Explicitly const: pins overload resolution to the copying
+        // Optimize(const ModelProto&, ...) even though `proto` happens to be
+        // unused afterward here, so this call site's behavior can't change
+        // silently if a consuming overload is ever added upstream of it (see
+        // Optimizer::optimize(ModelProto&, ...) in optimize.h).
+        auto const result = optimization::Optimize(
+            static_cast<const ModelProto&>(proto), names);
         std::string out;
         result.SerializeToString(&out);
         return nb::bytes(out.data(), out.size());
@@ -42,7 +48,8 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
       [](const nb::bytes& bytes, const std::vector<std::string>& names) {
         ModelProto proto{};
         ParseProtoFromPyBytes(&proto, bytes);
-        auto const result = optimization::OptimizeFixed(proto, names);
+        auto const result = optimization::OptimizeFixed(
+            static_cast<const ModelProto&>(proto), names);
         std::string out;
         result.SerializeToString(&out);
         return nb::bytes(out.data(), out.size());
@@ -55,7 +62,8 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
                                const std::string& export_data_file_name) {
         ModelProto proto{};
         optimization::loadModel(&proto, import_model_path, true);
-        auto result = optimization::Optimize(proto, names);
+        auto result = optimization::Optimize(
+            static_cast<const ModelProto&>(proto), names);
         optimization::saveModel(&result, export_model_path, true,
                                 export_data_file_name);
       });
@@ -68,13 +76,20 @@ NB_MODULE(onnx_opt_cpp2py_export, onnx_opt_cpp2py_export) {
          const std::string& export_data_file_name) {
         ModelProto proto{};
         optimization::loadModel(&proto, import_model_path, true);
-        auto result = optimization::OptimizeFixed(proto, names);
+        auto result = optimization::OptimizeFixed(
+            static_cast<const ModelProto&>(proto), names);
         optimization::saveModel(&result, export_model_path, true,
                                 export_data_file_name);
       });
   onnx_opt_cpp2py_export.def("get_available_passes",
                              &optimization::GetAvailablePasses);
   onnx_opt_cpp2py_export.def("get_fuse_and_elimination_passes",
                              &optimization::GetFuseAndEliminationPass);
+  // Toggle whether the passes treat graph initializers as constant tensors
+  // (default true). See SetInitializersAsConstants in optimize.h.
+  onnx_opt_cpp2py_export.def("set_initializers_as_constants",
+                             &optimization::SetInitializersAsConstants);
+  onnx_opt_cpp2py_export.def("initializers_as_constants",
+                             &optimization::InitializersAsConstants);
 }
 }  // namespace ONNX_NAMESPACE

onnxoptimizer/optimize.cc

@@ -31,16 +31,50 @@ 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, report);
+}
+void OptimizeGraph(
+    Graph& graph,
+    const std::vector<std::string>& names,
+    std::map<std::string, unsigned int>* report,
+    bool clear_tensor_digest_cache) {
+  Optimizer current_opt(names, false);
+  current_opt.optimize(graph, report, clear_tensor_digest_cache);
+}
+void OptimizeGraphFixed(
+    Graph& graph,
+    const std::vector<std::string>& names,
+    std::map<std::string, unsigned int>* report,
+    bool clear_tensor_digest_cache) {
+  Optimizer current_opt(names, true);
+  current_opt.optimize(graph, report, clear_tensor_digest_cache);
+}
+#ifdef ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS
+ModelProto Optimize(
+    ModelProto& mp_in,
+    const std::vector<std::string>& names,
+    std::map<std::string, unsigned int>* report) {
+  Optimizer current_opt(names, false);
+  return current_opt.optimize(mp_in, report);
+}
+ModelProto OptimizeFixed(
+    ModelProto& mp_in,
+    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);
 }
+#endif  // ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS
 const std::vector<std::string> GetAvailablePasses() {
   return Optimizer::passes.GetAvailablePasses();
 }

onnxoptimizer/optimize.h

@@ -10,10 +10,10 @@
 #include "onnx/common/ir.h"
 #include "onnx/common/ir_pb_converter.h"
 #include "onnx/proto_utils.h"
-
 #include "onnxoptimizer/pass_manager.h"
 #include "onnxoptimizer/pass_registry.h"
-
+#include "onnxoptimizer/passes/cse_util.h"
+#include "onnxoptimizer/passes/tensor_content_hash.h"
 #include "vector"
 
 namespace ONNX_NAMESPACE {
@@ -26,8 +26,51 @@ struct Optimizer {
   Optimizer(const std::vector<std::string> &names, const bool fixed_point);
   ~Optimizer();
 
-  ModelProto optimize(const ModelProto &_mp_in) {
-    const ModelProto* mp_in = &_mp_in;
+  // Optimize the ONNX C++ IR (Graph) in place, running the configured passes
+  // directly on the graph. This avoids the ModelProto <-> Graph round-trip
+  // entirely and is intended for C++ callers that already hold a Graph (e.g.
+  // onnxsim's OptAndShape fixed point, which imports once and can keep
+  // re-running passes on the same Graph across rounds where shape inference
+  // made no change -- see onnxsim issue #633). 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.
+  //
+  // 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,
+  // matching the ModelProto-based optimize() below.
+  //
+  // If `clear_tensor_digest_cache` is true (the default, and correct for
+  // essentially every caller), the two tensor-hash caches consulted by
+  // eliminate_duplicate_initializer and eliminate_common_subexpression --
+  // TensorContentDigest's (tensor_content_hash.h, the typed-field path) and
+  // CSETensorHash's raw_data-branch cache (cse_util.h's g_raw_hash_cache,
+  // the common path for real exported models) -- are cleared before
+  // running the passes, bounding their memory to the tensors this one
+  // optimize() call touches. Pass false only if the caller itself manages
+  // those caches' lifetime across several optimize() calls on the *same*
+  // resident Graph (e.g. onnxsim's OptAndShape fixed point, which calls
+  // this once per round but wants hashes computed in an earlier round to
+  // stay cached in a later one) -- see ClearTensorContentDigestCache's
+  // header comment for why that's safe to do explicitly (the same
+  // reasoning applies to ClearRawHashCache).
+  void optimize(Graph &graph,
+                std::map<std::string, unsigned int> *report = nullptr,
+                bool clear_tensor_digest_cache = true) {
+    if (clear_tensor_digest_cache) {
+      ClearTensorContentDigestCache();
+      ClearRawHashCache();
+    }
+    auto analysis = this->pass_manager->run(graph);
+    if (report != nullptr && analysis != nullptr) {
+      *report = analysis->transform_counts;
+    }
+  }
+
+  // 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) {
       // Upgrade ir_version to 4 so that initializer can be not in input
@@ -46,21 +89,68 @@ struct Optimizer {
     }
 
     ModelProto mp_out = PrepareOutput(*mp_in);
-    this->pass_manager->run(*g);
+    this->optimize(*g, report);
     ExportModelProto(&mp_out, g);
 
     // Maybe we can optimize these functions, now just copy
     AddFunctionsToModel(*mp_in, mp_out);
     return mp_out;
   }
 
+#ifdef ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS
+  // Consuming overload: same as above, but moves each initializer's raw
+  // bytes out of `mp_in` on Import and out of the internal Graph on Export,
+  // instead of copying them at each end of the ModelProto<->Graph round
+  // trip. This roughly halves the memory traffic of one optimize() call for
+  // weight-heavy models (see onnxsim issue #633), at the cost of leaving
+  // `mp_in`'s initializer tensors with empty raw data afterward. Only call
+  // this when `mp_in` is about to be discarded or overwritten by the caller
+  // -- e.g. onnxsim's OptAndShape fixed point, which immediately
+  // move-assigns this call's return value back over its input model on
+  // every iteration.
+  //
+  // Only defined when compiled against an onnx fork that provides the
+  // matching consuming ImportModelProto/ExportModelProto overloads (see
+  // ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS in ir_pb_converter.h) --
+  // e.g. absent when this library is linked against onnxruntime's own
+  // bundled, unpatched onnx copy instead.
+  ModelProto optimize(ModelProto &mp_in,
+                      std::map<std::string, unsigned int> *report = nullptr) {
+    if (mp_in.ir_version() == 3) {
+      // Rare legacy path; not worth threading the moving Import/Export
+      // through, so fall back to the copying overload above.
+      const ModelProto &const_mp_in = mp_in;
+      return optimize(const_mp_in, report);
+    }
+    std::shared_ptr<Graph> g(ImportModelProto(mp_in));
+
+    if (g.get() == nullptr) {
+      std::cerr << "Warning: onnx optimizer is unable to parse input model. "
+                << "(The IR version of the ONNX model may be too old.)"
+                << std::endl;
+      // If we can't parse the file, just return the input. ImportModelProto
+      // fails before touching any tensor data (it only checks ir_version),
+      // so mp_in is still intact here.
+      return mp_in;
+    }
+
+    ModelProto mp_out = PrepareOutput(mp_in);
+    this->optimize(*g, report);
+    ExportModelProto(&mp_out, g, /*consume_tensor_data=*/true);
+
+    // Maybe we can optimize these functions, now just copy
+    AddFunctionsToModel(mp_in, mp_out);
+    return mp_out;
+  }
+#endif  // ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS
+
  private:
   std::shared_ptr<PassManager> pass_manager;
 
   void AddFunctionsToModel(const ModelProto &original_model,
                            ModelProto &output_model) {
-    for (const auto& function_proto : original_model.functions()) {
-      auto* p_f = output_model.add_functions();
+    for (const auto &function_proto : original_model.functions()) {
+      auto *p_f = output_model.add_functions();
       p_f->CopyFrom(function_proto);
     }
   }
@@ -94,10 +184,55 @@ const std::vector<std::string> GetAvailablePasses();
 
 const std::vector<std::string> GetFuseAndEliminationPass();
 
+// Control whether the optimizer passes treat graph initializers as constant
+// tensors. The default (true) is onnxoptimizer's historical behaviour, in which
+// an initializer-backed value is a constant and value-baking passes
+// (fuse_bn_into_conv, fuse_add_bias_into_conv, nop-reshape/expand on a constant
+// shape, ...) may consume and fold it. When set to false, initializers are
+// treated as non-constant, so those passes leave initializer-backed values --
+// and the weights they represent -- untouched; Constant *nodes* are still
+// treated as constants. The setting is thread-local and stays in effect until
+// changed, so callers that flip it should restore it afterwards.
+void SetInitializersAsConstants(bool value);
+bool InitializersAsConstants();
+
 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);
+
+// In-place counterparts that operate directly on the ONNX C++ IR (Graph),
+// skipping the ModelProto <-> Graph conversion entirely. For C++ callers
+// that already hold a Graph -- see Optimizer::optimize(Graph&, ...)'s doc
+// comment. Unlike the consuming ModelProto overloads below, these do not
+// depend on ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS: they never touch
+// ModelProto at all, so they work identically whether this library is
+// linked against onnxsim's onnx fork or onnxruntime's bundled, unpatched
+// onnx copy.
+// `clear_tensor_digest_cache`: see Optimizer::optimize(Graph&, ...)'s doc
+// comment above -- the default (true) is correct for essentially every
+// caller.
+void OptimizeGraph(Graph &graph, const std::vector<std::string> &names,
+                   std::map<std::string, unsigned int> *report = nullptr,
+                   bool clear_tensor_digest_cache = true);
+
+void OptimizeGraphFixed(Graph &graph, const std::vector<std::string> &names,
+                        std::map<std::string, unsigned int> *report = nullptr,
+                        bool clear_tensor_digest_cache = true);
+
+#ifdef ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS
+// Consuming overloads: see Optimizer::optimize(ModelProto&, ...)'s doc
+// comment. Only call these when `mp_in` is about to be discarded or
+// overwritten by the caller.
+ModelProto Optimize(ModelProto &mp_in, const std::vector<std::string> &names,
+                    std::map<std::string, unsigned int> *report = nullptr);
+
+ModelProto OptimizeFixed(ModelProto &mp_in,
+                         const std::vector<std::string> &names,
+                         std::map<std::string, unsigned int> *report = nullptr);
+#endif  // ONNX_IR_PB_CONVERTER_HAS_CONSUMING_OVERLOADS
 }  // namespace optimization
 }  // namespace ONNX_NAMESPACE