Switch to Select optimisations - Two Case switch to select

[kariddi]

Summary

This optimization is derived from the optimisations in http://reviews.llvm.org/D5525

The optimisations in D5525 are of three kind. This patch contains only one kind (the Convert Two case switch to select), for review.

Checkout D5525 initial comment for a full description of all the optimisations.

[hans]

Thanks for breaking this out to a separate patch! It makes it much easier to review.

I think this is ready to be committed with the comments below addressed. Do you have commit access, or would you like me to land it for you?

lib/Transforms/Utils/SimplifyCFG.cpp

3453

nit: there is no need to repeat the function name in the comment. This goes for the other functions too.

And instead of "This is an helper function used to..." I'd just put "Helper function that adds CaseVal to the list of cases that generate Result."

3602

The lines from "const unsigned BitWidth" to "computeKnownBits(...)" look like leftovers from the other patch.

test/Transforms/SimplifyCFG/switch-to-select-two-case.ll

16

Please check the full icmp instruction, i.e. capture the name and what is being compared.

52

Please check the full icmp instruction.

[kariddi]

I do have commit access :)
Let me know if you think it is better removing function names from the comments or not considering most of the functions have them, and after addressing this last thing I will commit it.

Thanks again!

lib/Transforms/Utils/SimplifyCFG.cpp

3453

I put the name in the comment, because most of the other functions in the file follow that style.
I thought it was a good idea to follow the general style in the file.

As far as "This is ..." is concerned I think you are right, that piece is redundant.

3602

Whoops :D

test/Transforms/SimplifyCFG/switch-to-select-two-case.ll

16

Done

[kariddi]

test/Transforms/SimplifyCFG/switch-to-select-two-case.ll

52

done

[kariddi]

Latest comments taken care of.

[hans]

You're right, repeating the function name in the comment seems to be the dominating style in this file. Let's not worry about it, then.

[kariddi]

Landed r219223

lib/Transforms/Utils/SimplifyCFG.cpp

 STATISTIC(NumSpeculations, "Number of speculative executed instructions");
 
 namespace {
+  // The first field contains the value that the switch produces when a certain
+  // case group is selected, and the second field is a vector containing the cases
+  // composing the case group.
+  typedef SmallVector<std::pair<Constant *, SmallVector<ConstantInt *, 4>>, 2>
+    SwitchCaseResultVectorTy;
+  // The first field contains the phi node that generates a result of the switch
+  // and the second field contains the value generated for a certain case in the switch
+  // for that PHI.
+  typedef SmallVector<std::pair<PHINode *, Constant *>, 4> SwitchCaseResultsTy;
+
   /// ValueEqualityComparisonCase - Represents a case of a switch.
   struct ValueEqualityComparisonCase {
     ConstantInt *Value;

   return Res.size() > 0;
 }
 
+// MapCaseToResult - Helper function used to
+// add CaseVal to the list of cases that generate Result.
+static void MapCaseToResult(ConstantInt *CaseVal,
+    SwitchCaseResultVectorTy &UniqueResults,
+    Constant *Result) {
+  for (auto &I : UniqueResults) {
+    if (I.first == Result) {
+      I.second.push_back(CaseVal);
+      return;
+    }
+  }
+  UniqueResults.push_back(std::make_pair(Result,
+        SmallVector<ConstantInt*, 4>(1, CaseVal)));
+}
+
+// InitializeUniqueCases - 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, const DataLayout *DL, PHINode *&PHI,
+    BasicBlock *&CommonDest,
+    SwitchCaseResultVectorTy &UniqueResults,
+    Constant *&DefaultResult) {
+  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))
+      return false;
+
+    // Only one value per case is permitted
+    if (Results.size() > 1)
+      return false;
+    MapCaseToResult(CaseVal, UniqueResults, Results.begin()->second);
+
+    // 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;
+  BasicBlock *DefaultDest = SI->getDefaultDest();
+  GetCaseResults(SI, nullptr, SI->getDefaultDest(), &CommonDest, DefaultResults,
+                 DL);
+  // If the default value is not found abort unless the default destination
+  // is unreachable.
+  DefaultResult =
+      DefaultResults.size() == 1 ? DefaultResults.begin()->second : nullptr;
+  if ((!DefaultResult &&
+        !isa<UnreachableInst>(DefaultDest->getFirstNonPHIOrDbg())))
+    return false;
+
+  return true;
+}
+
+// ConvertTwoCaseSwitch - Helper function that checks if it is possible to
+// transform a switch with only two cases (or two cases + default)
+// that produces a result into a value select.
+// Example:
+// switch (a) {
+//   case 10:                %0 = icmp eq i32 %a, 10
+//     return 10;            %1 = select i1 %0, i32 10, i32 4
+//   case 20:        ---->   %2 = icmp eq i32 %a, 20
+//     return 2;             %3 = select i1 %2, i32 2, i32 %1
+//   default:
+//     return 4;
+// }
+static Value *
+ConvertTwoCaseSwitch(const SwitchCaseResultVectorTy &ResultVector,
+                     Constant *DefaultResult, Value *Condition,
+                     IRBuilder<> &Builder) {
+  assert(ResultVector.size() == 2 &&
+      "We should have exactly two unique results at this point");
+  // If we are selecting between only two cases transform into a simple
+  // select or a two-way select if default is possible.
+  if (ResultVector[0].second.size() == 1 &&
+      ResultVector[1].second.size() == 1) {
+    ConstantInt *const FirstCase = ResultVector[0].second[0];
+    ConstantInt *const SecondCase = ResultVector[1].second[0];
+
+    bool DefaultCanTrigger = DefaultResult;
+    Value *SelectValue = ResultVector[1].first;
+    if (DefaultCanTrigger) {
+      Value *const ValueCompare =
+          Builder.CreateICmpEQ(Condition, SecondCase, "switch.selectcmp");
+      SelectValue = Builder.CreateSelect(ValueCompare, ResultVector[1].first,
+                                         DefaultResult, "switch.select");
+    }
+    Value *const ValueCompare =
+        Builder.CreateICmpEQ(Condition, FirstCase, "switch.selectcmp");
+    return Builder.CreateSelect(ValueCompare, ResultVector[0].first, SelectValue,
+                                "switch.select");
+  }
+
+  return nullptr;
+}
+
+// RemoveSwitchAfterSelectConversion - Helper function to cleanup a switch
+// instruction that has been converted into a select, fixing up PHI nodes and
+// basic blocks.
+static void RemoveSwitchAfterSelectConversion(SwitchInst *SI, PHINode *PHI,
+                                              Value *SelectValue,
+                                              IRBuilder<> &Builder) {
+  BasicBlock *SelectBB = SI->getParent();
+  if (PHI->getBasicBlockIndex(SelectBB) >= 0)
+    PHI->removeIncomingValue(SelectBB);
+  PHI->addIncoming(SelectValue, SelectBB);
+
+  Builder.CreateBr(PHI->getParent());
+
+  // Remove the switch.
+  for (unsigned i = 0, e = SI->getNumSuccessors(); i < e; ++i) {
+    BasicBlock *Succ = SI->getSuccessor(i);
+
+    if (Succ == PHI->getParent())
+      continue;
+    Succ->removePredecessor(SelectBB);
+  }
+  SI->eraseFromParent();
+}
+
+/// SwitchToSelect - If the switch is only used to initialize one or more
+/// phi nodes in a common successor block with only two different
+/// constant values, replace the switch with select.
+static bool SwitchToSelect(SwitchInst *SI, IRBuilder<> &Builder,
+                           const DataLayout *DL, AssumptionTracker *AT) {
+  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, DL, PHI, CommonDest, UniqueResults,
+                             DefaultResult))
+    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 = ConvertTwoCaseSwitch(
+      UniqueResults,
+      DefaultResult, Cond, Builder);
+  if (SelectValue) {
+    RemoveSwitchAfterSelectConversion(SI, PHI, SelectValue, Builder);
+    return true;
+  }
+  // The switch couldn't be converted into a select.
+  return false;
+}
+
 namespace {
   /// SwitchLookupTable - This class represents a lookup table that can be used
   /// to replace a switch.

   if (EliminateDeadSwitchCases(SI, DL, AT))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, DL, AT) | true;
 
+  if (SwitchToSelect(SI, Builder, DL, AT))
+    return SimplifyCFG(BB, TTI, BonusInstThreshold, DL, AT) | true;
+
   if (ForwardSwitchConditionToPHI(SI))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, DL, AT) | true;
 

test/Transforms/SimplifyCFG/UnreachableEliminate.ll

 }
 
 ; PR9450
-define i32 @test4(i32 %v) {
+define i32 @test4(i32 %v, i32 %w) {
 ; CHECK: entry:
 ; CHECK-NEXT:  switch i32 %v, label %T [
 ; CHECK-NEXT:    i32 3, label %V

 default:
         unreachable
 U:
-        ret i32 1
+        ret i32 %w
 T:
         ret i32 2
 }

test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll

   %x = phi i32 [ 3, %sw.default ], [ 7, %sw.bb1 ], [ 9, %entry ]
   ret i32 %x
 ; CHECK-LABEL: @twocases(
-; CHECK: switch i32
+; CHECK-NOT: switch i32
 ; CHECK-NOT: @switch.table
+; CHECK: %switch.selectcmp
+; CHECK-NEXT: %switch.select
+; CHECK-NEXT: %switch.selectcmp1
+; CHECK-NEXT: %switch.select2
 }
 
 ; Don't build tables for switches with TLS variables.

test/Transforms/SimplifyCFG/switch-to-select-two-case.ll

+; RUN: opt < %s -simplifycfg -S | FileCheck %s
+
+; int foo1_with_default(int a) {
+;   switch(a) {
+;     case 10:
+;       return 10;
+;     case 20:
+;       return 2;
+;   }
+;   return 4;
+; }
+
+define i32 @foo1_with_default(i32 %a) {
+; CHECK-LABEL: @foo1_with_default
+; CHECK: %switch.selectcmp = icmp eq i32 %a, 20
+; CHECK-NEXT: %switch.select = select i1 %switch.selectcmp, i32 2, i32 4
+; CHECK-NEXT: %switch.selectcmp1 = icmp eq i32 %a, 10
+; CHECK-NEXT: %switch.select2 = select i1 %switch.selectcmp1, i32 10, i32 %switch.select
+entry:
+  switch i32 %a, label %sw.epilog [
+    i32 10, label %sw.bb
+    i32 20, label %sw.bb1
+  ]
+
+sw.bb:
+  br label %return
+
+sw.bb1:
+  br label %return
+
+sw.epilog:
+  br label %return
+
+return:
+  %retval.0 = phi i32 [ 4, %sw.epilog ], [ 2, %sw.bb1 ], [ 10, %sw.bb ]
+  ret i32 %retval.0
+}
+
+; int foo1_without_default(int a) {
+;   switch(a) {
+;     case 10:
+;       return 10;
+;     case 20:
+;       return 2;
+;   }
+;   __builtin_unreachable();
+; }
+
+define i32 @foo1_without_default(i32 %a) {
+; CHECK-LABEL: @foo1_without_default
+; CHECK: %switch.selectcmp = icmp eq i32 %a, 10
+; CHECK-NEXT: %switch.select = select i1 %switch.selectcmp, i32 10, i32 2
+; CHECK-NOT: %switch.selectcmp1
+entry:
+  switch i32 %a, label %sw.epilog [
+    i32 10, label %sw.bb
+    i32 20, label %sw.bb1
+  ]
+
+sw.bb:
+  br label %return
+
+sw.bb1:
+  br label %return
+
+sw.epilog:
+  unreachable
+
+return:
+  %retval.0 = phi i32 [ 2, %sw.bb1 ], [ 10, %sw.bb ]
+  ret i32 %retval.0
+}