[llvm] [Reassociate] Compute value ranks iteratively (PR #204952)
Hao Ren via llvm-commits
llvm-commits at lists.llvm.org
Mon Jun 29 15:55:23 PDT 2026
https://github.com/nvidia-moomoo updated https://github.com/llvm/llvm-project/pull/204952
>From d03b96ca4ffe20de3b1b5e1f8dd119283af70361 Mon Sep 17 00:00:00 2001
From: Hao Ren <haor at nvidia.com>
Date: Mon, 15 Jun 2026 23:03:53 +0000
Subject: [PATCH] [Reassociate] Compute value ranks iteratively
ReassociatePass::getRank recursively walks operand def-use chains before memoizing ranks. For a very deep acyclic chain whose tail feeds a reassociable operation, this can recurse once per chain element and overflow the native stack.
Compute ranks with an explicit post-order worklist instead. The worklist preserves the existing rank calculation and memoization behavior while bounding native stack use.
Add a regression test that generates a 65536-deep select chain and runs reassociate under an 8 MB stack.
---
llvm/lib/Transforms/Scalar/Reassociate.cpp | 82 +++++++++++++++----
.../Reassociate/Inputs/deep-select-chain.py | 19 +++++
.../Reassociate/deep-select-chain.ll | 17 ++++
3 files changed, 100 insertions(+), 18 deletions(-)
create mode 100644 llvm/test/Transforms/Reassociate/Inputs/deep-select-chain.py
create mode 100644 llvm/test/Transforms/Reassociate/deep-select-chain.ll
diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index a5105b84c3117..4cda6740526aa 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -212,24 +212,70 @@ unsigned ReassociatePass::getRank(Value *V) {
if (unsigned Rank = ValueRankMap[I])
return Rank; // Rank already known?
- // If this is an expression, return the 1+MAX(rank(LHS), rank(RHS)) so that
- // we can reassociate expressions for code motion! Since we do not recurse
- // for PHI nodes, we cannot have infinite recursion here, because there
- // cannot be loops in the value graph that do not go through PHI nodes.
- unsigned Rank = 0, MaxRank = RankMap[I->getParent()];
- for (unsigned i = 0, e = I->getNumOperands(); i != e && Rank != MaxRank; ++i)
- Rank = std::max(Rank, getRank(I->getOperand(i)));
-
- // If this is a 'not' or 'neg' instruction, do not count it for rank. This
- // assures us that X and ~X will have the same rank.
- if (!match(I, m_Not(m_Value())) && !match(I, m_Neg(m_Value())) &&
- !match(I, m_FNeg(m_Value())))
- ++Rank;
-
- LLVM_DEBUG(dbgs() << "Calculated Rank[" << V->getName() << "] = " << Rank
- << "\n");
-
- return ValueRankMap[I] = Rank;
+ // Return 1+MAX(rank(LHS), rank(RHS)) for expressions so we can reassociate
+ // expressions for code motion. Use an explicit worklist rather than native
+ // recursion so long acyclic use-def chains do not overflow the stack.
+ struct RankWorkItem {
+ Value *V;
+ unsigned OpNo;
+ unsigned Rank;
+ };
+
+ auto GetLeafRank = [this](Value *V) {
+ return isa<Argument>(V) ? ValueRankMap[V] : 0;
+ };
+
+ SmallVector<RankWorkItem, 16> Worklist;
+ // Each item is one suspended recursive getRank() call.
+ // Children are pushed first; completed ranks are folded back into the parent.
+ Worklist.push_back(RankWorkItem{I, 0, 0});
+
+ auto CompleteRank = [&](unsigned Rank) {
+ // Once the current use-def node has a known rank, carry that rank back to
+ // the parent expression and advance past the operand that led here.
+ Worklist.pop_back();
+ if (Worklist.empty())
+ return true;
+
+ RankWorkItem &Parent = Worklist.back();
+ Parent.Rank = std::max(Parent.Rank, Rank);
+ ++Parent.OpNo;
+ return false;
+ };
+
+ while (true) {
+ RankWorkItem &Item = Worklist.back();
+ Instruction *CurI = dyn_cast<Instruction>(Item.V);
+ unsigned Rank = 0;
+
+ if (!CurI) {
+ // Non-instruction value, such as a constant.
+ Rank = GetLeafRank(Item.V);
+ } else if (unsigned ExistingRank = ValueRankMap[CurI]) {
+ // Instruction that is not movable.
+ Rank = ExistingRank;
+ } else if (Item.OpNo == CurI->getNumOperands() ||
+ Item.Rank == RankMap[CurI->getParent()]) {
+ // All operands were visited or the max block rank was reached.
+ Rank = Item.Rank;
+ // If this is a 'not' or 'neg' instruction, do not count it for rank.
+ // This assures us that X and ~X will have the same rank.
+ if (!match(CurI, m_Not(m_Value())) && !match(CurI, m_Neg(m_Value())) &&
+ !match(CurI, m_FNeg(m_Value())))
+ ++Rank;
+
+ LLVM_DEBUG(dbgs() << "Calculated Rank[" << CurI->getName()
+ << "] = " << Rank << "\n");
+
+ ValueRankMap[CurI] = Rank;
+ } else {
+ Worklist.push_back(RankWorkItem{CurI->getOperand(Item.OpNo), 0, 0});
+ continue;
+ }
+
+ if (CompleteRank(Rank))
+ return Rank;
+ }
}
// Canonicalize constants to RHS. Otherwise, sort the operands by rank.
diff --git a/llvm/test/Transforms/Reassociate/Inputs/deep-select-chain.py b/llvm/test/Transforms/Reassociate/Inputs/deep-select-chain.py
new file mode 100644
index 0000000000000..637782e6f09a5
--- /dev/null
+++ b/llvm/test/Transforms/Reassociate/Inputs/deep-select-chain.py
@@ -0,0 +1,19 @@
+import sys
+
+
+def main():
+ depth = int(sys.argv[1])
+ assert depth > 0
+
+ print("define i32 @deep_select_chain(i1 %p, i32 %a, i32 %b) {")
+ print("entry:")
+ print(" %s0 = select i1 %p, i32 %a, i32 %b")
+ for i in range(1, depth):
+ print(f" %s{i} = select i1 %p, i32 %s{i - 1}, i32 %b")
+ print(f" %r = add i32 %s{depth - 1}, %a")
+ print(" ret i32 %r")
+ print("}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/llvm/test/Transforms/Reassociate/deep-select-chain.ll b/llvm/test/Transforms/Reassociate/deep-select-chain.ll
new file mode 100644
index 0000000000000..f32f4ffa39bc3
--- /dev/null
+++ b/llvm/test/Transforms/Reassociate/deep-select-chain.ll
@@ -0,0 +1,17 @@
+; RUN: opt -S -passes=reassociate < %s | FileCheck %s
+; RUN: %python %S/Inputs/deep-select-chain.py 65536 > %t.deep.ll
+; RUN: ulimit -s 8192 && opt -S -passes=reassociate < %t.deep.ll -o /dev/null
+; REQUIRES: system-linux
+
+define i32 @small_select_chain(i1 %p, i32 %a, i32 %b) {
+; CHECK-LABEL: @small_select_chain(
+; CHECK: %s0 = select i1 %p, i32 %a, i32 %b
+; CHECK-NEXT: %s1 = select i1 %p, i32 %s0, i32 %b
+; CHECK-NEXT: %r = add i32 %s1, %a
+; CHECK-NEXT: ret i32 %r
+;
+ %s0 = select i1 %p, i32 %a, i32 %b
+ %s1 = select i1 %p, i32 %s0, i32 %b
+ %r = add i32 %s1, %a
+ ret i32 %r
+}
More information about the llvm-commits
mailing list