Cache CSETensorHash's raw_data hash by tensor_id(), the actual bottleneck

onnxoptimizer/optimize.cc

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

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 {
@@ -38,8 +38,28 @@ struct Optimizer {
   // 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) {
+                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;
@@ -50,7 +70,7 @@ struct Optimizer {
   // 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;
+    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
@@ -129,8 +149,8 @@ struct Optimizer {
 
   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);
     }
   }
@@ -192,18 +212,22 @@ ModelProto OptimizeFixed(const ModelProto &mp_in,
 // 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);
+                   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);
+                        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,
+ModelProto Optimize(ModelProto &mp_in, const std::vector<std::string> &names,
                     std::map<std::string, unsigned int> *report = nullptr);
 
 ModelProto OptimizeFixed(ModelProto &mp_in,

onnxoptimizer/pass.cc

@@ -2,17 +2,92 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-#include "onnx/common/assertions.h"
-
 #include "onnxoptimizer/pass.h"
 
+#include <chrono>
+
+#include "onnx/common/assertions.h"
+
 namespace ONNX_NAMESPACE {
 namespace optimization {
 
-Pass::Pass(
-    PassType pass_type,
-    PassEfficiency pass_efficiency,
-    PassOptimizationType pass_optimization_type) {
+namespace {
+bool g_pass_phase_profiling_enabled = false;
+std::unordered_map<std::string, PassPhaseTiming> g_pass_phase_timings;
+std::unordered_map<std::string, PassTotalTiming> g_pass_total_timings;
+CSEPassTiming g_cse_pass_timing;
+DeadendPassTiming g_deadend_pass_timing;
+}  // namespace
+
+void SetPassPhaseProfilingEnabled(bool enabled) {
+  g_pass_phase_profiling_enabled = enabled;
+}
+
+bool GetPassPhaseProfilingEnabled() {
+  return g_pass_phase_profiling_enabled;
+}
+
+const std::unordered_map<std::string, PassPhaseTiming>& GetPassPhaseTimings() {
+  return g_pass_phase_timings;
+}
+
+void ResetPassPhaseTimings() {
+  g_pass_phase_timings.clear();
+}
+
+void RecordPassTotalTime(const std::string& pass_name, double ms) {
+  PassTotalTiming& t = g_pass_total_timings[pass_name];
+  t.calls++;
+  t.total_ms += ms;
+}
+
+const std::unordered_map<std::string, PassTotalTiming>& GetPassTotalTimings() {
+  return g_pass_total_timings;
+}
+
+void ResetPassTotalTimings() {
+  g_pass_total_timings.clear();
+}
+
+void RecordCSEPassTiming(uint64_t nodes_seen, uint64_t nodes_filtered_out,
+                         uint64_t nodes_replaced, double filter_ms,
+                         double lookup_ms, double replace_ms) {
+  g_cse_pass_timing.calls++;
+  g_cse_pass_timing.nodes_seen += nodes_seen;
+  g_cse_pass_timing.nodes_filtered_out += nodes_filtered_out;
+  g_cse_pass_timing.nodes_replaced += nodes_replaced;
+  g_cse_pass_timing.filter_ms += filter_ms;
+  g_cse_pass_timing.lookup_ms += lookup_ms;
+  g_cse_pass_timing.replace_ms += replace_ms;
+}
+
+const CSEPassTiming& GetCSEPassTiming() {
+  return g_cse_pass_timing;
+}
+
+void ResetCSEPassTiming() {
+  g_cse_pass_timing = CSEPassTiming();
+}
+
+void RecordDeadendPassTiming(uint64_t nodes_seen, uint64_t nodes_removed,
+                             double has_uses_ms, double destroy_ms) {
+  g_deadend_pass_timing.calls++;
+  g_deadend_pass_timing.nodes_seen += nodes_seen;
+  g_deadend_pass_timing.nodes_removed += nodes_removed;
+  g_deadend_pass_timing.has_uses_ms += has_uses_ms;
+  g_deadend_pass_timing.destroy_ms += destroy_ms;
+}
+
+const DeadendPassTiming& GetDeadendPassTiming() {
+  return g_deadend_pass_timing;
+}
+
+void ResetDeadendPassTiming() {
+  g_deadend_pass_timing = DeadendPassTiming();
+}
+
+Pass::Pass(PassType pass_type, PassEfficiency pass_efficiency,
+           PassOptimizationType pass_optimization_type) {
   this->pass_type = pass_type;
   this->pass_efficiency = pass_efficiency;
   this->pass_optimization_type = pass_optimization_type;
@@ -21,8 +96,7 @@ Pass::Pass(
 Pass::~Pass() {}
 
 unsigned int Pass::DescendOnGraphAttributesAndCount(
-    Node* n,
-    std::function<unsigned int(Graph&)> fn) {
+    Node* n, std::function<unsigned int(Graph&)> fn) {
   unsigned int num_changes = 0;
   for (auto name : n->attributeNames()) {
     auto kind = n->kindOf(name);
@@ -39,8 +113,7 @@ unsigned int Pass::DescendOnGraphAttributesAndCount(
 }
 
 void Pass::DescendOnGraphAttributesUnconstrained(
-    Node* n,
-    std::function<void(Graph&)> fn) {
+    Node* n, std::function<void(Graph&)> fn) {
   for (auto name : n->attributeNames()) {
     auto kind = n->kindOf(name);
     if (kind == AttributeKind::g) {
@@ -58,13 +131,39 @@ PredicateBasedPass::~PredicateBasedPass() {}
 
 unsigned int PredicateBasedPass::_runPassInternal(Graph& graph) {
   unsigned int num_changes = false;
+  // Only touches g_pass_phase_timings when profiling is on, so the lookup
+  // (once per call, not once per node) and the two std::chrono reads per
+  // node below are the only cost this diagnostic imposes when enabled.
+  const bool profiling = g_pass_phase_profiling_enabled;
+  PassPhaseTiming* timing =
+      profiling ? &g_pass_phase_timings[this->getPassName()] : nullptr;
   for (auto it = graph.begin(); it != graph.end(); ++it) {
     auto* n = *it;
     num_changes += this->DescendOnGraphAttributesAndCount(
         n, [this](Graph& g) { return _runPassInternal(g); });
-    if (this->patternMatchPredicate(n)) {
+    bool matched;
+    if (profiling) {
+      const auto t0 = std::chrono::steady_clock::now();
+      matched = this->patternMatchPredicate(n);
+      const auto t1 = std::chrono::steady_clock::now();
+      timing->match_calls++;
+      timing->match_ms +=
+          std::chrono::duration<double, std::milli>(t1 - t0).count();
+    } else {
+      matched = this->patternMatchPredicate(n);
+    }
+    if (matched) {
       NodeDestroyType destroy_type = NodeDestroyType::DestroyZero;
-      num_changes += this->runTransform(n, graph, destroy_type);
+      if (profiling) {
+        const auto t0 = std::chrono::steady_clock::now();
+        num_changes += this->runTransform(n, graph, destroy_type);
+        const auto t1 = std::chrono::steady_clock::now();
+        timing->transform_calls++;
+        timing->transform_ms +=
+            std::chrono::duration<double, std::milli>(t1 - t0).count();
+      } else {
+        num_changes += this->runTransform(n, graph, destroy_type);
+      }
 
       if (destroy_type == NodeDestroyType::DestroyOne) {
         it.destroyCurrent();
@@ -88,9 +187,7 @@ std::shared_ptr<PostPassAnalysis> PredicateBasedPass::runPass(Graph& graph) {
 }
 
 CountBasedPassAnalysis::CountBasedPassAnalysis(
-    Pass* pass,
-    unsigned int num_positive_transforms,
-    bool initialization_done,
+    Pass* pass, unsigned int num_positive_transforms, bool initialization_done,
     bool finalization_done) {
   this->pass = pass;
   this->num_positive_transforms = num_positive_transforms;
@@ -100,5 +197,5 @@ CountBasedPassAnalysis::CountBasedPassAnalysis(
 
 FullGraphBasedPass::~FullGraphBasedPass() {}
 
-} // namespace optimization
-} // namespace ONNX_NAMESPACE
+}  // namespace optimization
+}  // namespace ONNX_NAMESPACE

onnxoptimizer/pass.h

@@ -9,7 +9,10 @@
 
 #pragma once
 
+#include <cstdint>
 #include <string>
+#include <unordered_map>
+
 #include "onnx/common/ir.h"
 #include "onnx/onnx_pb.h"
 
@@ -21,6 +24,85 @@ struct PostPassAnalysis {
   virtual ~PostPassAnalysis() = default;
 };
 
+// Exploratory diagnostic: per-pass-name timing split between
+// PredicateBasedPass's two phases -- "matching" (patternMatchPredicate,
+// scanning nodes for rewrite candidates) and "modifying" (runTransform,
+// actually rewriting a matched node) -- written for onnxsim issue #633's
+// investigation into where OptimizeGraphFixed's ~50-round fixed point
+// actually spends its time. Covers PredicateBasedPass-derived passes only
+// (the majority of the default suite: fuse_*, most eliminate_*); the
+// smaller number of FullGraphBasedPass passes (eliminate_duplicate_
+// initializer, eliminate_common_subexpression, DCE, ...) implement their
+// own single-phase runPass() and aren't split by this.
+//
+// Off by default (SetPassPhaseProfilingEnabled(true) to turn on) so normal
+// runs pay zero std::chrono overhead. Not thread-safe to toggle
+// concurrently with a running pass, matching this library's other global
+// toggles (e.g. tensor_content_hash.h's SetTrustTensorContentHash).
+struct PassPhaseTiming {
+  uint64_t match_calls = 0;
+  double match_ms = 0.0;
+  uint64_t transform_calls = 0;
+  double transform_ms = 0.0;
+};
+void SetPassPhaseProfilingEnabled(bool enabled);
+bool GetPassPhaseProfilingEnabled();
+const std::unordered_map<std::string, PassPhaseTiming> &GetPassPhaseTimings();
+void ResetPassPhaseTimings();
+
+// Companion to PassPhaseTiming, at coarser granularity: total wall time
+// inside each pass's runPass(Graph&) call (FixedPointPassManager::run's
+// call sites), covering BOTH pass kinds uniformly -- PredicateBasedPass's
+// per-node loop overhead that PassPhaseTiming's match/transform timers don't
+// capture (iterator traversal, DescendOnGraphAttributesAndCount, ...) and
+// FullGraphBasedPass passes (eliminate_duplicate_initializer, eliminate_
+// common_subexpression, DCE, ...) that don't have a matching/modifying split
+// at all. Shares SetPassPhaseProfilingEnabled's on/off toggle.
+struct PassTotalTiming {
+  uint64_t calls = 0;
+  double total_ms = 0.0;
+};
+void RecordPassTotalTime(const std::string &pass_name, double ms);
+const std::unordered_map<std::string, PassTotalTiming> &GetPassTotalTimings();
+void ResetPassTotalTimings();
+
+// Internal breakdown of EliminateCommonSubexpressions's own per-node loop
+// (eliminate_common_subexpression.h), beyond what cse_util.h's CSENodeHash/
+// CSEEqual instrumentation already measures inside the hash-map lookup
+// itself. `lookup_ms` covers the whole `hash_map.emplace()` call (hashing
+// plus, on a bucket collision, CSEEqual), so it overlaps with cse_util.h's
+// node_hash_ms/node_equal_ms -- the two are complementary views of the same
+// work, not additive. Shares SetPassPhaseProfilingEnabled's on/off toggle.
+struct CSEPassTiming {
+  uint64_t calls = 0;
+  uint64_t nodes_seen = 0;
+  uint64_t nodes_filtered_out = 0;
+  uint64_t nodes_replaced = 0;
+  double filter_ms = 0.0;
+  double lookup_ms = 0.0;
+  double replace_ms = 0.0;
+};
+void RecordCSEPassTiming(uint64_t nodes_seen, uint64_t nodes_filtered_out,
+                         uint64_t nodes_replaced, double filter_ms,
+                         double lookup_ms, double replace_ms);
+const CSEPassTiming &GetCSEPassTiming();
+void ResetCSEPassTiming();
+
+// Internal breakdown of EliminateDead's own reverse-order sweep
+// (eliminate_deadend.h). Shares SetPassPhaseProfilingEnabled's on/off
+// toggle.
+struct DeadendPassTiming {
+  uint64_t calls = 0;
+  uint64_t nodes_seen = 0;
+  uint64_t nodes_removed = 0;
+  double has_uses_ms = 0.0;
+  double destroy_ms = 0.0;
+};
+void RecordDeadendPassTiming(uint64_t nodes_seen, uint64_t nodes_removed,
+                             double has_uses_ms, double destroy_ms);
+const DeadendPassTiming &GetDeadendPassTiming();
+void ResetDeadendPassTiming();
+
 // Enum that represents the type of optimization it is.
 enum PassType {
   // Class of optimizations that fuses operations.