From b4a750f2caf684db3db5c29794d6172c23d379a4 Mon Sep 17 00:00:00 2001 From: Dmitry Bushev Date: Mon, 23 Jun 2025 08:54:56 +0000 Subject: [PATCH 1/3] [DFAJumpThreading] Prevent pass from using too much memory. The limit 'dfa-max-num-paths' that is used to control number of enumerated paths was not checked against inside getPathsFromStateDefMap. It may lead to large memory consumption for complex enough switch statements. --- llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp index 938aab5879044..ac7af712bcf12 100644 --- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp +++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp @@ -618,6 +618,8 @@ struct AllSwitchPaths { VisitedBlocks UniqueBlocks; for (auto *IncomingBB : Phi->blocks()) { + if(Res.size() >= MaxNumPaths) + break; if (!UniqueBlocks.insert(IncomingBB).second) continue; if (!SwitchOuterLoop->contains(IncomingBB)) @@ -657,6 +659,8 @@ struct AllSwitchPaths { getPathsFromStateDefMap(StateDef, IncomingPhi, VB); for (ThreadingPath &Path : PredPaths) { Path.push_back(PhiBB); + if(Res.size() >= MaxNumPaths) + break; Res.push_back(std::move(Path)); } continue; @@ -679,6 +683,10 @@ struct AllSwitchPaths { ThreadingPath NewPath(Path); NewPath.appendExcludingFirst(IPath); NewPath.push_back(PhiBB); + if(Res.size() >= MaxNumPaths) { + VB.erase(PhiBB); + return Res; + } Res.push_back(NewPath); } } From 6cec630c8dd95ea106e4754ac9a4c7bc050287e8 Mon Sep 17 00:00:00 2001 From: Dmitry Bushev Date: Tue, 24 Jun 2025 09:46:39 +0000 Subject: [PATCH 2/3] Fix formatting --- llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp index ac7af712bcf12..54cf3c523bdff 100644 --- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp +++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp @@ -618,7 +618,7 @@ struct AllSwitchPaths { VisitedBlocks UniqueBlocks; for (auto *IncomingBB : Phi->blocks()) { - if(Res.size() >= MaxNumPaths) + if (Res.size() >= MaxNumPaths) break; if (!UniqueBlocks.insert(IncomingBB).second) continue; @@ -659,7 +659,7 @@ struct AllSwitchPaths { getPathsFromStateDefMap(StateDef, IncomingPhi, VB); for (ThreadingPath &Path : PredPaths) { Path.push_back(PhiBB); - if(Res.size() >= MaxNumPaths) + if (Res.size() >= MaxNumPaths) break; Res.push_back(std::move(Path)); } @@ -683,10 +683,10 @@ struct AllSwitchPaths { ThreadingPath NewPath(Path); NewPath.appendExcludingFirst(IPath); NewPath.push_back(PhiBB); - if(Res.size() >= MaxNumPaths) { + if (Res.size() >= MaxNumPaths) { VB.erase(PhiBB); return Res; - } + } Res.push_back(NewPath); } } From 2a13e447dde34fadcdd5ea0deac56c06e5f4e28c Mon Sep 17 00:00:00 2001 From: Dmitry Bushev Date: Tue, 22 Jul 2025 13:39:18 +0000 Subject: [PATCH 3/3] applied suggestion of @UsmanNadeem --- .../Transforms/Scalar/DFAJumpThreading.cpp | 44 ++++++++++--------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp index 54cf3c523bdff..a7ba54f03e61e 100644 --- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp +++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp @@ -582,15 +582,17 @@ struct AllSwitchPaths { VisitedBlocks VB; // Get paths from the determinator BBs to SwitchPhiDefBB std::vector PathsToPhiDef = - getPathsFromStateDefMap(StateDef, SwitchPhi, VB); + getPathsFromStateDefMap(StateDef, SwitchPhi, VB, MaxNumPaths); if (SwitchPhiDefBB == SwitchBlock) { TPaths = std::move(PathsToPhiDef); return; } + assert(MaxNumPaths >= PathsToPhiDef.size()); + auto PathsLimit = MaxNumPaths / PathsToPhiDef.size(); // Find and append paths from SwitchPhiDefBB to SwitchBlock. PathsType PathsToSwitchBB = - paths(SwitchPhiDefBB, SwitchBlock, VB, /* PathDepth = */ 1); + paths(SwitchPhiDefBB, SwitchBlock, VB, /* PathDepth = */ 1, PathsLimit); if (PathsToSwitchBB.empty()) return; @@ -611,14 +613,15 @@ struct AllSwitchPaths { typedef DenseMap StateDefMap; std::vector getPathsFromStateDefMap(StateDefMap &StateDef, PHINode *Phi, - VisitedBlocks &VB) { + VisitedBlocks &VB, + unsigned PathsLimit) { std::vector Res; auto *PhiBB = Phi->getParent(); VB.insert(PhiBB); VisitedBlocks UniqueBlocks; for (auto *IncomingBB : Phi->blocks()) { - if (Res.size() >= MaxNumPaths) + if (Res.size() >= PathsLimit) break; if (!UniqueBlocks.insert(IncomingBB).second) continue; @@ -655,12 +658,11 @@ struct AllSwitchPaths { // Direct predecessor, just add to the path. if (IncomingPhiDefBB == IncomingBB) { - std::vector PredPaths = - getPathsFromStateDefMap(StateDef, IncomingPhi, VB); + assert(PathsLimit > Res.size()); + std::vector PredPaths = getPathsFromStateDefMap( + StateDef, IncomingPhi, VB, PathsLimit - Res.size()); for (ThreadingPath &Path : PredPaths) { Path.push_back(PhiBB); - if (Res.size() >= MaxNumPaths) - break; Res.push_back(std::move(Path)); } continue; @@ -671,22 +673,22 @@ struct AllSwitchPaths { continue; PathsType IntermediatePaths; - IntermediatePaths = - paths(IncomingPhiDefBB, IncomingBB, VB, /* PathDepth = */ 1); + assert(PathsLimit > Res.size()); + auto InterPathLimit = PathsLimit - Res.size(); + IntermediatePaths = paths(IncomingPhiDefBB, IncomingBB, VB, + /* PathDepth = */ 1, InterPathLimit); if (IntermediatePaths.empty()) continue; + assert(InterPathLimit >= IntermediatePaths.size()); + auto PredPathLimit = InterPathLimit / IntermediatePaths.size(); std::vector PredPaths = - getPathsFromStateDefMap(StateDef, IncomingPhi, VB); + getPathsFromStateDefMap(StateDef, IncomingPhi, VB, PredPathLimit); for (const ThreadingPath &Path : PredPaths) { for (const PathType &IPath : IntermediatePaths) { ThreadingPath NewPath(Path); NewPath.appendExcludingFirst(IPath); NewPath.push_back(PhiBB); - if (Res.size() >= MaxNumPaths) { - VB.erase(PhiBB); - return Res; - } Res.push_back(NewPath); } } @@ -696,7 +698,7 @@ struct AllSwitchPaths { } PathsType paths(BasicBlock *BB, BasicBlock *ToBB, VisitedBlocks &Visited, - unsigned PathDepth) { + unsigned PathDepth, unsigned PathsLimit) { PathsType Res; // Stop exploring paths after visiting MaxPathLength blocks @@ -723,6 +725,8 @@ struct AllSwitchPaths { // is used to prevent a duplicate path from being generated SmallSet Successors; for (BasicBlock *Succ : successors(BB)) { + if (Res.size() >= PathsLimit) + break; if (!Successors.insert(Succ).second) continue; @@ -744,14 +748,12 @@ struct AllSwitchPaths { // coverage and compile time. if (LI->getLoopFor(Succ) != CurrLoop) continue; - - PathsType SuccPaths = paths(Succ, ToBB, Visited, PathDepth + 1); + assert(PathsLimit > Res.size()); + PathsType SuccPaths = + paths(Succ, ToBB, Visited, PathDepth + 1, PathsLimit - Res.size()); for (PathType &Path : SuccPaths) { Path.push_front(BB); Res.push_back(Path); - if (Res.size() >= MaxNumPaths) { - return Res; - } } } // This block could now be visited again from a different predecessor. Note