[SimplifyCFG] Improove and speed up ReduceSwitchRange

[shawnl]

Summary

I am trying to move towards much more space-efficient wwitch statements, using popcnt, as described in PR39013. This is the first patch towards that goal.

ReduceSwitchRange only ran if the switch statement wasn't "dense", resulting in unnecessary sparse switch statements. This is a very cheap transform, if it can make improvements should be run, so run it all the time. (that is all IsSwitchDense did, it did not determine if a jump table or lookup table should be generated) Doing this even when not generating a Lookup/Jump Table actually improves even LowerSwitch, as LowerSwitch wants the range to be constrained. Also:

• Do not do an expensive general GCD calculation when only countTrailingZeros() information is used from it.
• Re-order values to ideal with a subtraction, instead of only covering some cases by treating the values as signed.
• Implement the ctlz/cttz transform that was suggested in the comments. This transform is still not ideal, as SwitchToLookupTable doesn't know the potential range of values has been reduced, and there is no good plumbing in LLVM (that I know of) to get that information, and KnownBits is not sufficient.
• Do not reduce the switch range in SwitchToLookupTable

Gotcha: There a few cases where the old NeedMask would have been better than the sparse map, which patch I have yet to submit. However with only a few holes (8 bytes of holes) the old NeedMask implementation would use more space, as it doesn't use popcnt. Also, removing that feature means this will not be suitable for merging until it's replacement is approved.

[shawnl]

I realized that CommonRightDigits will always be 0 because of the Base subtraction. I will submit a patch without that confusing complexity.

[lebedev.ri]

I'd recommend git diff --name-only --cached | xargs -n 1 git blame --porcelain | grep "^author " | sort | uniq -c | sort -nr | head -10 to find reviewers.

[nikic]

lib/Transforms/Utils/SimplifyCFG.cpp

5579

Just looking at this, shouldn't there be some changes to existing tests?

[lebedev.ri]

lib/Transforms/Utils/SimplifyCFG.cpp

5498

So what about -Os/-Oz?

5629–5633

Did you see this comment?

[shawnl]

I think this patch will make more sense together with the other changes I want to make. I will withdraw it and submit the changes together.

lib/Transforms/Utils/SimplifyCFG.cpp

5498

This is not making the table, that calculation is elsewhere. This optimization is cheap and should always be applied.

5629–5633

This change does not change whether a lookup table gets created (although I am working on such patches). It only simplifies the switch.

[shawnl]

Add a second patch that changes SwitchToLookupTable

There were a number of cases where a table was smaller and we still didn't use it, and
also i8 tables are small enough that we should use them. Also, remove xor and and when producing a covered table.

[shawnl]

This is a combination of two patches. To see them separately: https://github.com/shawnl/llvm/commits/test

This patch now depends on D60823

lib/Transforms/Utils/SimplifyCFG.cpp

5579

There were no tests for the old behavior, surprisingly.

[shawnl]

get the clz transform working

[shawnl]

include context

[lebedev.ri]

Can you please write a human-readable description for the differential, in layman terms, that makes sense?
Rough bullet point list could be:
"Motivation. Issues with current implementations. The possible solutions. The implemented solution. Drawbacks, gotchas of the change.", etc.

This change is pretty big. It would be really best to do this in smaller steps.
Test coverage in particular is bothering me - this only has basic positive tests.
No negative tests. No tests for boundary conditions

lib/Transforms/Utils/SimplifyCFG.cpp

5131

What is this 24, 3*8? Make it variable with a name.

5177

This sounds like a hardcoded value for some specific target.

5323–5326

Why is this being changed? Can this be a separate patch?

5392–5393

Constants are always on RHS anyway

test/Transforms/SimplifyCFG/switch-simplify-range.ll

3

Please don't use -O? (or for that matter, don't passes other than the pass from directory name)
in tests outside of PassOrdering directory.

[shawnl]

add comment that range analysis would be useful for the clz/ctz transform

[shawnl]

just the first patch

[shawnl]

lib/Transforms/Utils/SimplifyCFG.cpp

5323–5326

ReduceSwitchRange now re-orders values so that they start at zero, so is superior to this simple "sort by signed" trick.

[shawnl]

Add a comment clarifying that what looks like a edge condition is not.

[shawnl]

add another comment about avoiding an edge condition.

[shawnl]

no need to change this

[shawnl]

Remove some useless code now that we actually use ReduceSwitchRange(). However I did not bother to update NeedMask, because the whole plan is to implement something much more general than it.

[shawnl]

remove use of -O2 in tests

[jmolloy]

Hi Shawn,

Thanks for being patient waiting for me to review this. You mention that there are no tests for this behaviour; that's very strange as I remember committing this change with tests . In fact, it was test/Transforms/SimplifyCFG/rangereduce.ll (https://github.com/llvm/llvm-project/commit/b2e436de423aac1de4d764166317dd7ab9cbf751). I'm surprised this test didn't need updating.

Please reupload with full context.

My general comment on this patch is that it's too busy with too few comments to give a good review on. For changes like this I'd really appreciate splitting it into functionally distinct parts. For example here we have:

1. SwitchLookupTable no longer takes an offset and assumes things start from zero and can use an unsigned comparison. There is clearly some change in ReduceSwitchRange that makes this true but I can't unpick the relevant hunk from the rest of this patch. Note, I'm not yet convinced this is correct but the review is too cluttered for me to really understand the logic change.
2. The removal of GCD and use of clz (on APInt) for GCD calculation.
3. The use of clz/cttz
4. Change shift/shift/or to fshr, which I'm certain is going to screw up some backend.

These should really be in separate patches.

Cheers,

James

lib/Transforms/Utils/SimplifyCFG.cpp

5295

I'm having trouble groking at a glance why this is correct. Could you comment? Previously we had a slt/sgt pair, now just a ugt?

Note that if you'd have uploaded with full context I could go look closer to understand more ;)

5389–5390

This can't be committed as-is. If you can't modify the existing code to work in the interim between your two patches, the other patch should be merged into this one.

Note this will make it more of a pain to review, so I'd suggest trying hard to modify the existing code :)

5466

I'm having trouble working out how this heuristic relates to the existing code. Have you changed the heuristic as well as the generated code?

Note, a NFC change is much easier to review and get landed.

5482

You're describing a tricky gotcha here. Please be more specific and go into more detail as to why it's there in the first place.

5486

This whole calculation should be outlined into a helper function. Then for example you wouldn't need to nest the whole thing inside if (BitWidth > 8).

5495

use std::min/std::max here?

5500

Please change the phrasing. Instead of describing what we might do *without* the check, please describe what the goal of the check is (negative vs. positive sense makes it easier to grok what the condition is and whether it matches the descriptive comment).

5501

might be worse

... but why?

5502

Why > 64?

5504

This whole comment has me worried. Target-specific guessing games don't belong in SimplifyCFG :)

A heuristic that isn't obviously target-agnostic should get a TTI hook. They're cheap, go for it! :)

5518

You're not only mutating but rearranging Values while iterating? this is an antipattern. Please remove.

5584

Why not use Builder.CreateCall here?

5594

How well supported is codegen for fshr? shift/shift/or for rotate will go through pretty much any backend; I've been out of the loop for a little while - how new is fshr?

[shawnl]

Split into multiple patches, starting at D61150

lib/Transforms/Utils/SimplifyCFG.cpp

 class SwitchLookupTable {
 public:
   /// Create a lookup table to use as a switch replacement with the contents
   /// of Values, using DefaultValue to fill any holes in the table.
   SwitchLookupTable(
-      Module &M, uint64_t TableSize, ConstantInt *Offset,
+      Module &M, uint64_t TableSize,
       const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
       Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName);
 
   /// Build instructions with Builder to retrieve the value at
   /// the position given by Index in the lookup table.

 };
 
 } // end anonymous namespace
 
 SwitchLookupTable::SwitchLookupTable(
-    Module &M, uint64_t TableSize, ConstantInt *Offset,
+    Module &M, uint64_t TableSize,
     const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
     Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName) {
   assert(Values.size() && "Can't build lookup table without values!");
   assert(TableSize >= Values.size() && "Can't fit values in table!");
 

   for (size_t I = 0, E = Values.size(); I != E; ++I) {
     ConstantInt *CaseVal = Values[I].first;
     Constant *CaseRes = Values[I].second;
     assert(CaseRes->getType() == ValueType);
 
-    uint64_t Idx = (CaseVal->getValue() - Offset->getValue()).getLimitedValue();
+    uint64_t Idx = CaseVal->getValue().getLimitedValue();
     TableContents[Idx] = CaseRes;
 
     if (CaseRes != SingleValue)
       SingleValue = nullptr;
   }

 
   // Figure out the corresponding result for each case value and phi node in the
   // common destination, as well as the min and max case values.
   assert(!empty(SI->cases()));
   SwitchInst::CaseIt CI = SI->case_begin();
-  ConstantInt *MinCaseVal = CI->getCaseValue();
   ConstantInt *MaxCaseVal = CI->getCaseValue();
 
   BasicBlock *CommonDest = nullptr;
 
   using ResultListTy = SmallVector<std::pair<ConstantInt *, Constant *>, 4>;

   SmallDenseMap<PHINode *, Type *> ResultTypes;
   SmallVector<PHINode *, 4> PHIs;
 
   for (SwitchInst::CaseIt E = SI->case_end(); CI != E; ++CI) {
     ConstantInt *CaseVal = CI->getCaseValue();
-    if (CaseVal->getValue().slt(MinCaseVal->getValue()))
-      MinCaseVal = CaseVal;
-    if (CaseVal->getValue().sgt(MaxCaseVal->getValue()))
+    if (CaseVal->getValue().ugt(MaxCaseVal->getValue()))
       MaxCaseVal = CaseVal;
 
     // Resulting value at phi nodes for this case value.
     using ResultsTy = SmallVector<std::pair<PHINode *, Constant *>, 4>;
     ResultsTy Results;

   for (PHINode *PHI : PHIs) {
     ResultTypes[PHI] = ResultLists[PHI][0].second->getType();
   }
 
   uint64_t NumResults = ResultLists[PHIs[0]].size();
-  APInt RangeSpread = MaxCaseVal->getValue() - MinCaseVal->getValue();
-  uint64_t TableSize = RangeSpread.getLimitedValue() + 1;
+  uint64_t TableSize = MaxCaseVal->getValue().getLimitedValue() + 1;
   bool TableHasHoles = (NumResults < TableSize);
 
   // If the table has holes, we need a constant result for the default case
   // or a bitmask that fits in a register.
   SmallVector<std::pair<PHINode *, Constant *>, 4> DefaultResultsList;

   BasicBlock *LookupBB = BasicBlock::Create(
       Mod.getContext(), "switch.lookup", CommonDest->getParent(), CommonDest);
 
   // Compute the table index value.
   Builder.SetInsertPoint(SI);
-  Value *TableIndex;
-  if (MinCaseVal->isNullValue())
-    TableIndex = SI->getCondition();
-  else
-    TableIndex = Builder.CreateSub(SI->getCondition(), MinCaseVal,
-                                   "switch.tableidx");
+  Value *TableIndex = SI->getCondition();
 
   // Compute the maximum table size representable by the integer type we are
   // switching upon.
-  unsigned CaseSize = MinCaseVal->getType()->getPrimitiveSizeInBits();
+  unsigned CaseSize = MaxCaseVal->getType()->getPrimitiveSizeInBits();
   uint64_t MaxTableSize = CaseSize > 63 ? UINT64_MAX : 1ULL << CaseSize;
   assert(MaxTableSize >= TableSize &&
          "It is impossible for a switch to have more entries than the max "
          "representable value of its input integer type's size.");
 

     Builder.CreateBr(LookupBB);
     // Note: We call removeProdecessor later since we need to be able to get the
     // PHI value for the default case in case we're using a bit mask.
   } else {
     Value *Cmp = Builder.CreateICmpULT(
-        TableIndex, ConstantInt::get(MinCaseVal->getType(), TableSize));
+        TableIndex, ConstantInt::get(MaxCaseVal->getType(), TableSize));
     RangeCheckBranch =
         Builder.CreateCondBr(Cmp, LookupBB, SI->getDefaultDest());
   }
 
   // Populate the BB that does the lookups.
   Builder.SetInsertPoint(LookupBB);
 
   if (NeedMask) {
-    // Before doing the lookup, we do the hole check. The LookupBB is therefore
-    // re-purposed to do the hole check, and we create a new LookupBB.
-    BasicBlock *MaskBB = LookupBB;
-    MaskBB->setName("switch.hole_check");
-    LookupBB = BasicBlock::Create(Mod.getContext(), "switch.lookup",
-                                  CommonDest->getParent(), CommonDest);
-
-    // Make the mask's bitwidth at least 8-bit and a power-of-2 to avoid
-    // unnecessary illegal types.
-    uint64_t TableSizePowOf2 = NextPowerOf2(std::max(7ULL, TableSize - 1ULL));
-    APInt MaskInt(TableSizePowOf2, 0);
-    APInt One(TableSizePowOf2, 1);
-    // Build bitmask; fill in a 1 bit for every case.
-    const ResultListTy &ResultList = ResultLists[PHIs[0]];
-    for (size_t I = 0, E = ResultList.size(); I != E; ++I) {
-      uint64_t Idx = (ResultList[I].first->getValue() - MinCaseVal->getValue())
-                         .getLimitedValue();
-      MaskInt |= One << Idx;
-    }
-    ConstantInt *TableMask = ConstantInt::get(Mod.getContext(), MaskInt);
-
-    // Get the TableIndex'th bit of the bitmask.
-    // If this bit is 0 (meaning hole) jump to the default destination,
-    // else continue with table lookup.
-    IntegerType *MapTy = TableMask->getType();
-    Value *MaskIndex =
-        Builder.CreateZExtOrTrunc(TableIndex, MapTy, "switch.maskindex");
-    Value *Shifted = Builder.CreateLShr(TableMask, MaskIndex, "switch.shifted");
-    Value *LoBit = Builder.CreateTrunc(
-        Shifted, Type::getInt1Ty(Mod.getContext()), "switch.lobit");
-    Builder.CreateCondBr(LoBit, LookupBB, SI->getDefaultDest());
-
-    Builder.SetInsertPoint(LookupBB);
-    AddPredecessorToBlock(SI->getDefaultDest(), MaskBB, SI->getParent());
+    // Re-written in a later patch
+    return false;
   }
 
   if (!DefaultIsReachable || GeneratingCoveredLookupTable) {
     // We cached PHINodes in PHIs. To avoid accessing deleted PHINodes later,
     // do not delete PHINodes here.

     const ResultListTy &ResultList = ResultLists[PHI];
 
     // If using a bitmask, use any value to fill the lookup table holes.
     Constant *DV = NeedMask ? ResultLists[PHI][0].second : DefaultResults[PHI];
     StringRef FuncName = Fn->getName();
-    SwitchLookupTable Table(Mod, TableSize, MinCaseVal, ResultList, DV, DL,
+    SwitchLookupTable Table(Mod, TableSize, ResultList, DV, DL,
                             FuncName);
 
     Value *Result = Table.BuildLookup(TableIndex, Builder);
 
     // If the result is used to return immediately from the function, we want to

   if (NeedMask)
     ++NumLookupTablesHoles;
   return true;
 }
 
-static bool isSwitchDense(ArrayRef<int64_t> Values) {
-  // See also SelectionDAGBuilder::isDense(), which this function was based on.
-  uint64_t Diff = (uint64_t)Values.back() - (uint64_t)Values.front();
-  uint64_t Range = Diff + 1;
-  uint64_t NumCases = Values.size();
-  // 40% is the default density for building a jump table in optsize/minsize mode.
-  uint64_t MinDensity = 40;
-
-  return NumCases * 100 >= Range * MinDensity;
-}
-
 /// Try to transform a switch that has "holes" in it to a contiguous sequence
 /// of cases.
 ///
 /// A switch such as: switch(i) {case 5: case 9: case 13: case 17:} can be
 /// range-reduced to: switch ((i-5) / 4) {case 0: case 1: case 2: case 3:}.

 /// This converts a sparse switch into a dense switch which allows better
 /// lowering and could also allow transforming into a lookup table.
 static bool ReduceSwitchRange(SwitchInst *SI, IRBuilder<> &Builder,
                               const DataLayout &DL,
                               const TargetTransformInfo &TTI) {
+  // The number of cases that need to be removed by a subtraction operation
+  // to make it worth using.
+  const unsigned SubThreshold = (SI->getFunction()->hasOptSize() ? 2 : 8);
+  bool MadeChanges = false;
   auto *CondTy = cast<IntegerType>(SI->getCondition()->getType());
-  if (CondTy->getIntegerBitWidth() > 64 ||
-      !DL.fitsInLegalInteger(CondTy->getIntegerBitWidth()))
-    return false;
-  // Only bother with this optimization if there are more than 3 switch cases;
-  // SDAG will only bother creating jump tables for 4 or more cases.
-  if (SI->getNumCases() < 4)
+  unsigned BitWidth = CondTy->getIntegerBitWidth();
+  if (BitWidth > 64 ||
+      !DL.fitsInLegalInteger(BitWidth))
     return false;
 
-  // This transform is agnostic to the signedness of the input or case values. We
-  // can treat the case values as signed or unsigned. We can optimize more common
-  // cases such as a sequence crossing zero {-4,0,4,8} if we interpret case values
-  // as signed.
-  SmallVector<int64_t,4> Values;
+  SmallVector<uint64_t,4> Values;
   for (auto &C : SI->cases())
-    Values.push_back(C.getCaseValue()->getValue().getSExtValue());
+    Values.push_back(C.getCaseValue()->getLimitedValue());
   llvm::sort(Values);
 
-  // If the switch is already dense, there's nothing useful to do here.
-  if (isSwitchDense(Values))
-    return false;
+  // A cheap speculative transform: handle power-of-two flags using
+  // @clz or @ctz as the key function.
+  // TODO: this transform would benifit from proper range analysis in SwitchToLookupTable
+  // BitWidth > 8 could be relaxed after this is fixed, but not eliminated, because
+  // this transforms 0 => BitWidth.
+  bool UseClz = false;
+  bool UseCtz = false;
+  if (BitWidth > 8) {
+    unsigned MaxPopCount = 0;
+    unsigned MinClz = 64, MaxClz = 0, MinCtz = 64, MaxCtz = 0;
+    for (auto &V : Values) {
+      MaxPopCount = countPopulation(V);
+      if (MaxPopCount > 1)
+        break;
+      unsigned Clz = APInt(BitWidth, V).countLeadingZeros();
+      unsigned Ctz = APInt(BitWidth, V).countTrailingZeros();
+      if (Clz < MinClz) MinClz = Clz;
+      if (Clz > MaxClz) MaxClz = Clz;
+      if (Ctz < MinCtz) MinCtz = Ctz;
+      if (Ctz > MaxCtz) MaxCtz = Ctz;
+    }
+    // Without this check we might do clz followed by ctz, and if Values contains 0,
+    // the result might be woorse.
+    if (MaxPopCount == 1 && Values.back() > 64) {
+      MadeChanges = true;
+      // Prefer clz because it is one instruction cheaper
+      // on ARM, but they cost the same on x86 so if we only need
+      // the subtraction on clz use ctz.
+      UseClz = (MinClz < SubThreshold) || (MinCtz >= SubThreshold);
+      if (UseClz)
+        for (auto &V : Values)
+          V = APInt(BitWidth, V).countLeadingZeros();
+      else {
+        UseCtz = true;
+        for (auto &V : Values) {
+          V = APInt(BitWidth, V).countTrailingZeros();
+        }
+      }
+      // 0 will suddenly become the largest (BitWidth), so we need to sort again.
+      llvm::sort(Values);
+    }
+  }
 
-  // First, transform the values such that they start at zero and ascend.
-  int64_t Base = Values[0];
-  for (auto &V : Values)
-    V -= (uint64_t)(Base);
+  // Find the element that has the most distance from it's previous, wrapping around.
+  uint64_t BestDistance = APInt::getMaxValue(CondTy->getIntegerBitWidth()).getLimitedValue() -
+    Values.back() + Values.front() + 1;
+  unsigned BestIndex = 0;
+  for (unsigned i = 1;i != Values.size();i++) {
+    if (Values[i] - Values[i-1] > BestDistance) {
+      BestIndex = i;
+      BestDistance = Values[i] - Values[i-1];
+    }
+  }
 
-  // Now we have signed numbers that have been shifted so that, given enough
-  // precision, there are no negative values. Since the rest of the transform
-  // is bitwise only, we switch now to an unsigned representation.
-  uint64_t GCD = 0;
-  for (auto &V : Values)
-    GCD = GreatestCommonDivisor64(GCD, (uint64_t)V);
+  uint64_t Base = 0;
+  // Now transform the values such that they start at zero and ascend.
+  if ((BestDistance > SubThreshold) &&
+    (BestIndex != 0 || (Values[0] >= SubThreshold))) {
+    Base = Values[BestIndex];
+    MadeChanges = true;
+    for (auto &V : Values)
+      V = (APInt(BitWidth, V) - Base).getLimitedValue();
+  }
 
   // This transform can be done speculatively because it is so cheap - it results
-  // in a single rotate operation being inserted. This can only happen if the
-  // factor extracted is a power of 2.
-  // FIXME: If the GCD is an odd number we can multiply by the multiplicative
-  // inverse of GCD and then perform this transform.
-  // FIXME: It's possible that optimizing a switch on powers of two might also
-  // be beneficial - flag values are often powers of two and we could use a CLZ
-  // as the key function.
-  if (GCD <= 1 || !isPowerOf2_64(GCD))
-    // No common divisor found or too expensive to compute key function.
-    return false;
-
-  unsigned Shift = Log2_64(GCD);
-  for (auto &V : Values)
-    V = (int64_t)((uint64_t)V >> Shift);
+  // in a single rotate operation being inserted.
+  unsigned Shift = 64;
+  for (auto &V : Values) {
+    // There is no edge condition when the BitWidth is less than 64, because if
+    // 0 is the only value then a shift does nothing, and LLVM requires
+    // well-formed IR to not have duplicate cases.
+    unsigned TZ = countTrailingZeros(V);
+    if (TZ < Shift) {
+      Shift = TZ;
+      if (Shift == 0)
+        break;
+    }
+  }
+  if (Shift) {
+    MadeChanges = true;
+    for (auto &V : Values)
+      V = V >> Shift;
+  }
 
-  if (!isSwitchDense(Values))
-    // Transform didn't create a dense switch.
+  if (!MadeChanges)
+    // Didn't do anything
     return false;
 
   // The obvious transform is to shift the switch condition right and emit a
   // check that the condition actually cleanly divided by GCD, i.e.
   //   C & (1 << Shift - 1) == 0

   // shift and puts the shifted-off bits in the uppermost bits. If any of these
   // are nonzero then the switch condition will be very large and will hit the
   // default case.
 
   auto *Ty = cast<IntegerType>(SI->getCondition()->getType());
-  Builder.SetInsertPoint(SI);
-  auto *ShiftC = ConstantInt::get(Ty, Shift);
-  auto *Sub = Builder.CreateSub(SI->getCondition(), ConstantInt::get(Ty, Base));
-  auto *LShr = Builder.CreateLShr(Sub, ShiftC);
-  auto *Shl = Builder.CreateShl(Sub, Ty->getBitWidth() - Shift);
-  auto *Rot = Builder.CreateOr(LShr, Shl);
-  SI->replaceUsesOfWith(SI->getCondition(), Rot);
+  {
+    auto Zero = ConstantInt::get(IntegerType::get(Ty->getContext(), 1), 0);
+    Builder.SetInsertPoint(SI);
+    Value *ZerosTransform;
+    if (UseClz) {
+      Function *Ctlz = Intrinsic::getDeclaration(SI->getModule(), Intrinsic::ctlz, Ty);
+      ZerosTransform = Builder.Insert(CallInst::Create(Ctlz, {SI->getCondition(), Zero}));
+    } else if (UseCtz) {
+      Function *Cttz = Intrinsic::getDeclaration(SI->getModule(), Intrinsic::cttz, Ty);
+      ZerosTransform = Builder.Insert(CallInst::Create(Cttz, {SI->getCondition(), Zero}));
+    } else
+      ZerosTransform = SI->getCondition();
+
+    auto *Sub = Builder.CreateSub(ZerosTransform, ConstantInt::get(Ty, Base));
+    Value *Key;
+    if (Shift) {
+      Function *Fshr = Intrinsic::getDeclaration(SI->getModule(), Intrinsic::fshr, Ty);
+      auto *ShiftC = ConstantInt::get(Ty, Shift);
+      Key = Builder.Insert(CallInst::Create(Fshr, {Sub, Sub, ShiftC}));
+    } else
+      Key = Sub;
+    SI->replaceUsesOfWith(SI->getCondition(), Key);
+  }
 
   for (auto Case : SI->cases()) {
     auto *Orig = Case.getCaseValue();
-    auto Sub = Orig->getValue() - APInt(Ty->getBitWidth(), Base);
+    uint64_t Zeros;
+    if (UseClz)
+      Zeros = Orig->getValue().countLeadingZeros();
+    else if (UseCtz)
+      Zeros = Orig->getValue().countTrailingZeros();
+    else
+      Zeros = Orig->getValue().getLimitedValue();
+    auto Sub = (APInt(BitWidth, Zeros) - Base).getLimitedValue();
     Case.setValue(
-        cast<ConstantInt>(ConstantInt::get(Ty, Sub.lshr(ShiftC->getValue()))));
+        cast<ConstantInt>(ConstantInt::get(Ty, APInt(BitWidth, Sub >> Shift))));
   }
+
   return true;
 }
 
 bool SimplifyCFGOpt::SimplifySwitch(SwitchInst *SI, IRBuilder<> &Builder) {
   BasicBlock *BB = SI->getParent();

     return requestResimplify();
 
   if (Options.ForwardSwitchCondToPhi && ForwardSwitchConditionToPHI(SI))
     return requestResimplify();
 
+  if (ReduceSwitchRange(SI, Builder, DL, TTI))
+    return requestResimplify();
+
   // The conversion from switch to lookup tables results in difficult-to-analyze
   // code and makes pruning branches much harder. This is a problem if the
   // switch expression itself can still be restricted as a result of inlining or
   // CVP. Therefore, only apply this transformation during late stages of the
   // optimisation pipeline.
   if (Options.ConvertSwitchToLookupTable &&
       SwitchToLookupTable(SI, Builder, DL, TTI))
     return requestResimplify();
 
-  if (ReduceSwitchRange(SI, Builder, DL, TTI))
-    return requestResimplify();
-
   return false;
 }
 
 bool SimplifyCFGOpt::SimplifyIndirectBr(IndirectBrInst *IBI) {
   BasicBlock *BB = IBI->getParent();

test/Transforms/SimplifyCFG/switch-simplify-range.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; Will not generate lookup tables without -O2 and the target triple below
+; RUN: opt -S -simplifycfg -O2 < %s | FileCheck %s
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+attributes #0 = { "no-jump-tables"="false" }
+
+define i64 @switch_common_right_bits(i8 %a) #0  {
+; CHECK-LABEL: @switch_common_right_bits(
+; CHECK-NEXT:  Entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i8 [[A:%.*]], 3
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i8 @llvm.fshl.i8(i8 [[TMP0]], i8 [[TMP0]], i8 7)
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i8 [[TMP1]], 4
+; CHECK-NEXT:    br i1 [[TMP2]], label [[SWITCH_LOOKUP:%.*]], label [[SWITCHELSE:%.*]]
+; CHECK:       switch.lookup:
+; CHECK-NEXT:    [[TMP3:%.*]] = sext i8 [[TMP1]] to i64
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [4 x i64], [4 x i64]* @switch.table.switch_common_right_bits, i64 0, i64 [[TMP3]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i64, i64* [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    ret i64 [[SWITCH_LOAD]]
+; CHECK:       SwitchElse:
+; CHECK-NEXT:    ret i64 10
+;
+Entry:
+  switch i8 %a, label %SwitchElse [
+  i8 253, label %SwitchProng
+  i8 255, label %SwitchProng1
+  i8 1, label %SwitchProng2
+  i8 3, label %SwitchProng3
+  ]
+SwitchElse:                                       ; preds = %Entry
+  ret i64 10
+SwitchProng:                                      ; preds = %Entry
+  ret i64 6
+SwitchProng1:                                     ; preds = %Entry
+  ret i64 3
+SwitchProng2:                                     ; preds = %Entry
+  ret i64 3
+SwitchProng3:                                     ; preds = %Entry
+  ret i64 3
+}
+
+define i64 @switch_ctz(i16 %a) optsize #0  {
+; CHECK-LABEL: @switch_ctz(
+; CHECK-NEXT:  Entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call i16 @llvm.cttz.i16(i16 [[A:%.*]], i1 false), !range !0
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i16 [[TMP0]], 8
+; CHECK-NEXT:    br i1 [[TMP1]], label [[SWITCH_LOOKUP:%.*]], label [[SWITCHELSE:%.*]]
+; CHECK:       switch.lookup:
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i16 [[TMP0]] to i64
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [8 x i64], [8 x i64]* @switch.table.switch_ctz, i64 0, i64 [[TMP2]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i64, i64* [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    ret i64 [[SWITCH_LOAD]]
+; CHECK:       SwitchElse:
+; CHECK-NEXT:    ret i64 10
+;
+Entry:
+  switch i16 %a, label %SwitchElse [
+  i16 2, label %SwitchProng
+  i16 4, label %SwitchProng1
+  i16 8, label %SwitchProng2
+  i16 1, label %SwitchProng3
+  i16 64, label %SwitchProng6
+  i16 128, label %SwitchProng7
+  i16 16, label %SwitchProng5
+  i16 32, label %SwitchProng4
+  ]
+SwitchElse:                                       ; preds = %Entry
+  ret i64 10
+SwitchProng:                                      ; preds = %Entry
+  ret i64 6
+SwitchProng1:                                     ; preds = %Entry
+  ret i64 3
+SwitchProng2:                                     ; preds = %Entry
+  ret i64 35
+SwitchProng3:                                     ; preds = %Entry
+  ret i64 31
+SwitchProng4:                                     ; preds = %Entry
+  ret i64 53
+SwitchProng5:                                     ; preds = %Entry
+  ret i64 51
+SwitchProng6:                                     ; preds = %Entry
+  ret i64 41
+SwitchProng7:                                     ; preds = %Entry
+  ret i64 34
+}
+
+define i64 @switch_clz(i8 %a) optsize #0  {
+; CHECK-LABEL: @switch_clz(
+; CHECK-NEXT:  Entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call i8 @llvm.fshl.i8(i8 [[A:%.*]], i8 [[A]], i8 3)
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i8 [[TMP0]], 5
+; CHECK-NEXT:    br i1 [[TMP1]], label [[SWITCH_LOOKUP:%.*]], label [[SWITCHPRONG2:%.*]]
+; CHECK:       switch.lookup:
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i8 [[TMP0]] to i64
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i64], [5 x i64]* @switch.table.switch_clz, i64 0, i64 [[TMP2]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i64, i64* [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    ret i64 [[SWITCH_LOAD]]
+; CHECK:       SwitchProng2:
+; CHECK-NEXT:    ret i64 12
+;
+Entry:
+  switch i8 %a, label %SwitchElse [
+  i8 128, label %SwitchProng2
+  i8 64, label %SwitchProng3
+  i8 32, label %SwitchProng6
+  i8 0, label %SwitchProng5
+  ]
+SwitchProng2:                                     ; preds = %Entry
+  ret i64 35
+SwitchProng3:                                     ; preds = %Entry
+  ret i64 31
+SwitchProng6:                                     ; preds = %Entry
+  ret i64 41
+SwitchProng5:                                     ; preds = %Entry
+  ret i64 40
+SwitchElse:                                     ; preds = %Entry
+  ret i64 12
+}
+
+;Must check that the default was filled in at index 0 as happened here:
+;@switch.table.switch_not_normalized_to_start_at_zero = private unnamed_addr constant [6 x i16] [i16 10, i16 7, i16 3, i16 1, i16 6, i16 8], align 2
+define i16 @switch_not_normalized_to_start_at_zero(i8 %a) #0  {
+; CHECK-LABEL: @switch_not_normalized_to_start_at_zero(
+; CHECK-NEXT:  Entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i8 [[A:%.*]], 6
+; CHECK-NEXT:    br i1 [[TMP0]], label [[SWITCH_LOOKUP:%.*]], label [[SWITCHELSE:%.*]]
+; CHECK:       switch.lookup:
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[A]] to i64
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [6 x i16], [6 x i16]* @switch.table.switch_not_normalized_to_start_at_zero, i64 0, i64 [[TMP1]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i16, i16* [[SWITCH_GEP]], align 2
+; CHECK-NEXT:    ret i16 [[SWITCH_LOAD]]
+; CHECK:       SwitchElse:
+; CHECK-NEXT:    ret i16 10
+;
+Entry:
+  switch i8 %a, label %SwitchElse [
+  i8 4, label %SwitchProng
+  i8 2, label %SwitchProng1
+  i8 1, label %SwitchProng2
+  i8 3, label %SwitchProng3
+  i8 5, label %SwitchProng4
+  ]
+SwitchElse:                                       ; preds = %Entry
+  ret i16 10
+SwitchProng:                                      ; preds = %Entry
+  ret i16 6
+SwitchProng1:                                     ; preds = %Entry
+  ret i16 3
+SwitchProng2:                                     ; preds = %Entry
+  ret i16 7
+SwitchProng3:                                     ; preds = %Entry
+  ret i16 1
+SwitchProng4:                                     ; preds = %Entry
+  ret i16 8
+}