[DFAJumpThreading] Prevent pass from using too much memory.

[dybv-sc]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Bushev Dmitry (dybv-sc)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/145482.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp (+8)

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);
         }
       }

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[dybv-sc]

Example of this pass using too much memory:

Compilation of function Perl_sv_vcatpvfn_flags of module sv.c from 600.perlbench_s benchmark with -mllvm --enable-dfa-jump-thread uses up to 16GB of heap. Almost all of this memory is allocated for threading paths. With the fix applied memory consumption shrinks down to less than ~500 MB.

[UsmanNadeem]

IntermediatePaths are appended to the list of PredPaths (here, and also in the run() function), so the total number of paths is the multiplication of the sizes of the two lists. IntermediatePaths are already calculated so we are wasting compute cycles here by doing wasteful computation and by breaking the loop later than we should.

Since getPathsFromStateDefMap() and paths() are called multiple times and from different places they don't see the global picture of the number of paths we already have.

My suggestion is to add an argument to these two functions and update every call to these functions with a more appropriate limit instead of MaxNumPaths.

[dybv-sc]

Great suggestion! Added path limit as parameter for both functions.

[dybv-sc]

@UsmanNadeem, @XChy, @nikic, could you please review this again?

[XChy]

LGTM

[evodius96]

Hi -- after this change, I am seeing a crash rooted this pass when building C++ runtimes in our downstream compiler (in particular, libcxx/src/locale.cpp). The crash happens at this line:

  `auto PathsLimit = MaxNumPaths / PathsToPhiDef.size();`                        

and PathsToPhiDef.size() ends up being 0 in this case.

I'm not sure this code handles all possible cases. Can you resolve or revert?

[dybv-sc]

@evodius96, Hi!
Thanks for your report, I've prepared new pull request with this case handled: #153193