From 66dff996a6514648d53a5fef6203d1e6211f5d04 Mon Sep 17 00:00:00 2001 From: Changqing Jing Date: Tue, 14 Apr 2026 12:27:07 +0800 Subject: [PATCH 1/3] avoid O(N^2) exiting-branch checks in CodeFolding --- src/passes/CodeFolding.cpp | 66 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 3 deletions(-) diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp index e53cd4f880d..e4e65a55940 100644 --- a/src/passes/CodeFolding.cpp +++ b/src/passes/CodeFolding.cpp @@ -63,6 +63,7 @@ #include "ir/effects.h" #include "ir/eh-utils.h" #include "ir/find_all.h" +#include "ir/iteration.h" #include "ir/label-utils.h" #include "ir/utils.h" #include "pass.h" @@ -135,6 +136,10 @@ struct CodeFolding modifieds; // modified code should not be processed // again, wait for next pass + // Cache for hasExitingBranches results. Populated by a single bottom-up + // walk to avoid O(N^2) repeated tree traversals on nested blocks. + std::unordered_map exitingBranchCache_; + // walking void visitExpression(Expression* curr) { @@ -299,6 +304,7 @@ struct CodeFolding returnTails.clear(); unoptimizables.clear(); modifieds.clear(); + exitingBranchCache_.clear(); if (needEHFixups) { EHUtils::handleBlockNestedPops(func, *getModule()); } @@ -306,6 +312,57 @@ struct CodeFolding } private: + // Pre-populate the exiting branch cache for all sub-expressions of root + // in a single O(N) bottom-up walk. After this, exitingBranchCache_ + // lookups are O(1). + void populateExitingBranchCache(Expression* root) { + struct CachePopulator + : public PostWalker> { + std::unordered_map& cache; + // Track unresolved branch targets at each node. We propagate children's + // targets upward: add uses, remove defs. If any remain, the expression + // has exiting branches. + std::unordered_map> targetSets; + + CachePopulator(std::unordered_map& cache) + : cache(cache) {} + + void visitExpression(Expression* curr) { + std::unordered_set targets; + // Merge children's target sets into ours (move to avoid copies) + ChildIterator children(curr); + for (auto* child : children) { + auto it = targetSets.find(child); + if (it != targetSets.end()) { + if (targets.empty()) { + targets = std::move(it->second); + } else { + targets.merge(it->second); + } + targetSets.erase(it); + } + } + // Add branch uses (names this expression branches to) + BranchUtils::operateOnScopeNameUses( + curr, [&](Name& name) { targets.insert(name); }); + // Remove branch defs (names this expression defines as targets) + BranchUtils::operateOnScopeNameDefs(curr, [&](Name& name) { + if (name.is()) { + targets.erase(name); + } + }); + bool hasExiting = !targets.empty(); + cache[curr] = hasExiting; + if (hasExiting) { + targetSets[curr] = std::move(targets); + } + } + }; + CachePopulator populator(exitingBranchCache_); + populator.walk(root); + } + // check if we can move a list of items out of another item. we can't do so // if one of the items has a branch to something inside outOf that is not // inside that item @@ -549,6 +606,11 @@ struct CodeFolding if (tails.size() < 2) { return false; } + // Pre-populate the cache once at the top level so all subsequent + // exitingBranchCache_ lookups are O(1). + if (num == 0) { + populateExitingBranchCache(getFunction()->body); + } // remove things that are untoward and cannot be optimized tails.erase( std::remove_if(tails.begin(), @@ -637,9 +699,7 @@ struct CodeFolding // TODO: this should not be a problem in // *non*-terminating tails, but // double-verify that - if (EffectAnalyzer( - getPassOptions(), *getModule(), newItem) - .hasExternalBreakTargets()) { + if (exitingBranchCache_[newItem]) { return true; } return false; From f263f08af5116d4b6198836e5e6604a6c43550fe Mon Sep 17 00:00:00 2001 From: Changqing Jing Date: Wed, 15 Apr 2026 12:41:17 +0800 Subject: [PATCH 2/3] Fix review --- src/passes/CodeFolding.cpp | 33 ++++++++++++++++++++------------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp index e4e65a55940..8bda2fc51db 100644 --- a/src/passes/CodeFolding.cpp +++ b/src/passes/CodeFolding.cpp @@ -136,9 +136,10 @@ struct CodeFolding modifieds; // modified code should not be processed // again, wait for next pass - // Cache for hasExitingBranches results. Populated by a single bottom-up - // walk to avoid O(N^2) repeated tree traversals on nested blocks. - std::unordered_map exitingBranchCache_; + // Cache of expressions that have branches exiting to targets defined + // outside them. Populated lazily on first access via PostWalker. + std::unordered_set exitingBranchCache_; + bool exitingBranchCachePopulated_ = false; // walking @@ -305,6 +306,7 @@ struct CodeFolding unoptimizables.clear(); modifieds.clear(); exitingBranchCache_.clear(); + exitingBranchCachePopulated_ = false; if (needEHFixups) { EHUtils::handleBlockNestedPops(func, *getModule()); } @@ -312,6 +314,17 @@ struct CodeFolding } private: + // Check if an expression has branches that exit to targets defined outside + // it. The cache is populated lazily on first call using a PostWalker for + // efficient bottom-up traversal. + bool hasExitingBranches(Expression* expr) { + if (!exitingBranchCachePopulated_) { + populateExitingBranchCache(getFunction()->body); + exitingBranchCachePopulated_ = true; + } + return exitingBranchCache_.count(expr); + } + // Pre-populate the exiting branch cache for all sub-expressions of root // in a single O(N) bottom-up walk. After this, exitingBranchCache_ // lookups are O(1). @@ -319,14 +332,13 @@ struct CodeFolding struct CachePopulator : public PostWalker> { - std::unordered_map& cache; + std::unordered_set& cache; // Track unresolved branch targets at each node. We propagate children's // targets upward: add uses, remove defs. If any remain, the expression // has exiting branches. std::unordered_map> targetSets; - CachePopulator(std::unordered_map& cache) - : cache(cache) {} + CachePopulator(std::unordered_set& cache) : cache(cache) {} void visitExpression(Expression* curr) { std::unordered_set targets; @@ -353,8 +365,8 @@ struct CodeFolding } }); bool hasExiting = !targets.empty(); - cache[curr] = hasExiting; if (hasExiting) { + cache.insert(curr); targetSets[curr] = std::move(targets); } } @@ -606,11 +618,6 @@ struct CodeFolding if (tails.size() < 2) { return false; } - // Pre-populate the cache once at the top level so all subsequent - // exitingBranchCache_ lookups are O(1). - if (num == 0) { - populateExitingBranchCache(getFunction()->body); - } // remove things that are untoward and cannot be optimized tails.erase( std::remove_if(tails.begin(), @@ -699,7 +706,7 @@ struct CodeFolding // TODO: this should not be a problem in // *non*-terminating tails, but // double-verify that - if (exitingBranchCache_[newItem]) { + if (hasExitingBranches(newItem)) { return true; } return false; From b90aee7cd2df94a40f13ae8555a3a638bfbd34f8 Mon Sep 17 00:00:00 2001 From: Changqing Jing Date: Thu, 16 Apr 2026 09:42:47 +0800 Subject: [PATCH 3/3] Fix review --- src/passes/CodeFolding.cpp | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp index 8bda2fc51db..20a8597ac11 100644 --- a/src/passes/CodeFolding.cpp +++ b/src/passes/CodeFolding.cpp @@ -136,11 +136,6 @@ struct CodeFolding modifieds; // modified code should not be processed // again, wait for next pass - // Cache of expressions that have branches exiting to targets defined - // outside them. Populated lazily on first access via PostWalker. - std::unordered_set exitingBranchCache_; - bool exitingBranchCachePopulated_ = false; - // walking void visitExpression(Expression* curr) { @@ -305,8 +300,8 @@ struct CodeFolding returnTails.clear(); unoptimizables.clear(); modifieds.clear(); - exitingBranchCache_.clear(); - exitingBranchCachePopulated_ = false; + exitingBranchCache.clear(); + exitingBranchCachePopulated = false; if (needEHFixups) { EHUtils::handleBlockNestedPops(func, *getModule()); } @@ -314,19 +309,21 @@ struct CodeFolding } private: - // Check if an expression has branches that exit to targets defined outside - // it. The cache is populated lazily on first call using a PostWalker for - // efficient bottom-up traversal. + // Cache of expressions that have branches exiting to targets defined + // outside them. Populated lazily on first access via hasExitingBranches(). + std::unordered_set exitingBranchCache; + bool exitingBranchCachePopulated = false; + bool hasExitingBranches(Expression* expr) { - if (!exitingBranchCachePopulated_) { + if (!exitingBranchCachePopulated) { populateExitingBranchCache(getFunction()->body); - exitingBranchCachePopulated_ = true; + exitingBranchCachePopulated = true; } - return exitingBranchCache_.count(expr); + return exitingBranchCache.count(expr); } // Pre-populate the exiting branch cache for all sub-expressions of root - // in a single O(N) bottom-up walk. After this, exitingBranchCache_ + // in a single O(N) bottom-up walk. After this, exitingBranchCache // lookups are O(1). void populateExitingBranchCache(Expression* root) { struct CachePopulator @@ -371,7 +368,7 @@ struct CodeFolding } } }; - CachePopulator populator(exitingBranchCache_); + CachePopulator populator(exitingBranchCache); populator.walk(root); }