Add early out to O(n^2) switch analysis in switch-to-select conversion

[AndrewScheidecker]

Summary

This adds an early out to the switch->select conversion analysis as soon as we've seen more than two unique results or cases. The conversion ultimately aborts in this case, but currently only does so after fully analyzing the switch, which takes O(NumCases^2) time if the switch is in a form that doesn't early out for other reasons.

In my WebAssembly VM that uses LLVM to generate machine code, I was running into this on a test case with a large, but simple, switch from the WebAssembly reference interpreter test suite: https://github.com/WebAssembly/spec/blob/master/test/core/br_table.wast#L110

To find the bottleneck, I increased the number of cases in the switch statement further, which caused it to spend 10 seconds in LLVM optimization passes. Adding this early out reduced that to tens of milliseconds.

[AndrewScheidecker]

Include full context in diff, as suggested on IRC

[kariddi]

Hi Andrew, your change modifies the semantics of the function, in the sense that it becomes specialized into just collecting the data to perform "ConvertTwoCaseSwitch()" and it will not be possible to use it for other potential optimization that require the same data.

Maybe we could add configurable cut-off points instead (in the form of default parameters) that default to what "ConvertTwoCaseSwitch()" requires for now (that would be 2 and 2)

Otherwise if we think that this function is not gonna be useful for any other optimization we could rename it and maybe it could even be simplified and the comment before it then should be updated to clarify the new use of it.

[AndrewScheidecker]

I added parameters to the InitializeUniqueCases function to control the early-outs, so it could theoretically be called by another transform that can support more complex switch statements.

[kariddi]

LGTM

[AndrewScheidecker]

I don't have commit access. Can somebody please commit this for me?

[kariddi]

Pushed as r322248

I forgot to mention that the patch was made by you in the original commit, so I made an extra commit (r322249) that mentions that the patch is actually yours :(

Sorry about that.

lib/Transforms/Utils/SimplifyCFG.cpp

for (PHINode &PHI : (*CommonDest)->phis()) {
 
     Res.push_back(std::make_pair(&PHI, ConstVal));
   }
 
   return Res.size() > 0;
 }
 
 // Helper function used to add CaseVal to the list of cases that generate
-// Result.
-static void MapCaseToResult(ConstantInt *CaseVal,
+// Result. Returns the updated number of cases that generate this result.
+static uintptr_t MapCaseToResult(ConstantInt *CaseVal,
                                  SwitchCaseResultVectorTy &UniqueResults,
                                  Constant *Result) {
   for (auto &I : UniqueResults) {
     if (I.first == Result) {
       I.second.push_back(CaseVal);
-      return;
+      return I.second.size();
     }
   }
   UniqueResults.push_back(
       std::make_pair(Result, SmallVector<ConstantInt *, 4>(1, CaseVal)));
+  return 1;
 }
 
 // Helper function that initializes a map containing
 // results for the PHI node of the common destination block for a switch
 // instruction. Returns false if multiple PHI nodes have been found or if
 // there is not a common destination block for the switch.
-static bool InitializeUniqueCases(SwitchInst *SI, PHINode *&PHI,
-                                  BasicBlock *&CommonDest,
-                                  SwitchCaseResultVectorTy &UniqueResults,
-                                  Constant *&DefaultResult,
-                                  const DataLayout &DL,
-                                  const TargetTransformInfo &TTI) {
+static bool
+InitializeUniqueCases(SwitchInst *SI, PHINode *&PHI, BasicBlock *&CommonDest,
+                      Constant *&DefaultResult, const DataLayout &DL,
+                      const TargetTransformInfo &TTI,
+                      uintptr_t MaxUniqueResults, uintptr_t MaxCasesPerResult) {
   for (auto &I : SI->cases()) {
     ConstantInt *CaseVal = I.getCaseValue();
 
     // Resulting value at phi nodes for this case value.
     SwitchCaseResultsTy Results;
     if (!GetCaseResults(SI, CaseVal, I.getCaseSuccessor(), &CommonDest, Results,
                         DL, TTI))
       return false;
 
     // Only one value per case is permitted
     if (Results.size() > 1)
       return false;
+
+    // Add the case->result mapping to UniqueResults.
+    const uintptr_t NumCasesForResult =
         MapCaseToResult(CaseVal, UniqueResults, Results.begin()->second);
 
+    // Early out if there are too many cases for this result.
+    if (NumCasesForResult > MaxCasesPerResult)
+      return false;
+
+    // Early out if there are too many unique results.
+    if (UniqueResults.size() > MaxUniqueResults)
+      return false;
+
     // Check the PHI consistency.
     if (!PHI)
       PHI = Results[0].first;
     else if (PHI != Results[0].first)
       return false;
   }
   // Find the default result value.
   SmallVector<std::pair<PHINode *, Constant *>, 1> DefaultResults;
static bool switchToSelect(SwitchInst *SI, IRBuilder<> &Builder,
                            const TargetTransformInfo &TTI) {
   Value *const Cond = SI->getCondition();
   PHINode *PHI = nullptr;
   BasicBlock *CommonDest = nullptr;
   Constant *DefaultResult;
   SwitchCaseResultVectorTy UniqueResults;
   // Collect all the cases that will deliver the same value from the switch.
   if (!InitializeUniqueCases(SI, PHI, CommonDest, UniqueResults, DefaultResult,
-                             DL, TTI))
+                             DL, TTI, 2, 1))
     return false;
   // Selects choose between maximum two values.
   if (UniqueResults.size() != 2)
     return false;
   assert(PHI != nullptr && "PHI for value select not found");
 
   Builder.SetInsertPoint(SI);
   Value *SelectValue =