[SimplifyCFG] Range reduce switches

[jmolloy]

Summary

If a switch is sparse and all the cases (once sorted) are in arithmetic progression, we can extract the common factor out of the switch and create a dense switch. For example:

switch (i) {
case 5: ...
case 9: ...
case 13: ...
case 17: ...
}

can become:

if ( (i - 5) % 4 ) goto default;
switch ((i - 5) / 4) {
case 0: ...
case 1: ...
case 2: ...
case 3: ...
}

The division and remainder operations could be costly so we only do this if the factor is a power of two. Dense switches can be lowered significantly better than sparse switches and can even be transformed into lookup tables.

[sbaranga]

lib/Transforms/Utils/SimplifyCFG.cpp

4993

Could you also add a test for this case?

5009

If it doesn't have holes then the GCD that you are looking for should be Values[1].

[jmolloy]

Hi Silviu,

Thanks for the review!

James

lib/Transforms/Utils/SimplifyCFG.cpp

4993

This should be @test3

5009

Uhhhhh, yes. Yes it should. Thanks for spotting this!

I've rewritten the code entirely.

[sanjoy]

lib/Transforms/Utils/SimplifyCFG.cpp

4989

I'd prefer a getZExtValue or getSExtValue here. That way the code will assert if someone later accidentally removed the check on the bit width.

4992

How about {-2, 0, 2, 4}? That's sparse and can benefit from this pass, but here you'll conclude that it has no holes.

5006

Don't you need to check that V is divisible by Divisor here?

5007

Might be worth calling out the signedness of the division here explicitly. I'd mildly prefer bitwise ops instead of division.

5017

I'd be tempted to emit the bitwise ops here directly. That way you don't do extra work later and are less sensitive to pass ordering.

[sanjoy]

lib/Transforms/Utils/SimplifyCFG.cpp

4992

Bad example, I meant cases like {-3, -1, 1, 3}.

[eli.friedman]

Maybe it makes sense to perform this sort of transformation in SelectionDAGBuilder::visitSwitch instead? Most of the relevant logic is there.

This patch only handles situations where all the values form a linear series, but it seems important to also catch cases where *most* of the values form a linear series (for example, 97, 105, 109, 113, 117, 121).

I think your profitability model is a bit off: we never generate a jump table for less than four cases, so trying to improve the density of a nonexistent jump table seems like a bad idea. And for cases where the divisor is small, the tradeoff between extra jump table entries and an extra compare+branch isn't obvious.

[jmolloy]

Hi Sanjoy and Eli,

Thanks for your reviews. I agree with all of your comments. This new version has a real density function and will perform the optimization if the switch is not dense to begin with and would be made dense. The density function is adapted from SelectionDAGBuilder, and I'm not too happy on replicating the heuristic here but I couldn't think of a better way.

Similarly the hardcoded bailout for < 4 cases - ideally I'd use TargetLowering here which has a hook, but we only have TTI. It might be worth adding a hook, but perhaps it's not so urgent so I've added a FIXME.

Eli, the reason I'm doing this in SimplifyCFG is because it's an enabler for switch->lookup table lowering (which is also in SimplifyCFG).

The testing has also been improved to cover more negative and unsigned large values and wraparound cases.

Cheers,

James

[eli.friedman]

lib/Transforms/Utils/SimplifyCFG.cpp

5022

Nit: please make the implicit casting here explicit.

5031

If the GCD is some odd number N, you can multiply by the inverse of N and switch on that. Not sure if that's actually useful in practice.

5056

If we were doing this in SelectionDAG, we would be able to combine this branch with the jump table's bounds check: "ROTATE(X - Base, Shift) > Limit".

[sanjoy]

Haven't done a thorough review yet, but some minor nitpicky comments inline.

lib/Transforms/Utils/SimplifyCFG.cpp

5029

This gcd computation looks like extra work -- why not:

uint64_t PotentialGCD = Values[1];
if (!isPowerOf2(PotentialGCD)) return false;
if (!all_of(Values, [](uint64_t V) { return V & (PotentialGCD - 1) == 0; }) return false;

5072

This linear search (via replaceUsesOfWith) seems wasteful. Why not change getCaseValue to return a Use & and do C.getCaseValue().set(New) or something similar?

[jmolloy]

lib/Transforms/Utils/SimplifyCFG.cpp

5056

Hmm, you're close to convincing me. I'm loath to give up the lookup table optimization though.

How gross (/ acceptable) would it be to implement this in SelectionDAG and then also teach lookup table lowering in SimplifyCFG this trick too?

[eli.friedman]

lib/Transforms/Utils/SimplifyCFG.cpp

5029

That isn't equivalent for something like "0, 8, 12, 16, 20".

5056

It wouldn't be too terrible as long as the actual algorithm for finding a reducible switch is factored out into a utility, I think.

[sanjoy]

lib/Transforms/Utils/SimplifyCFG.cpp

5029

Yeah, you're right. I was still thinking in terms of the older "100% dense" scheme.

[jmolloy]

Hi Eli,

That rotate trick is, quite simply, brilliant.

New version uses SelectionDAG - I'll factor some of the machinery out at a later date when I implement this for lookup table formation too.

All comments should now be addressed.

Cheers,

James

[eli.friedman]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

8332 ↗

(On Diff #60973)

Do we need to merge clusters in some cases, so we don't end up with adjacent clusters with the same destination? (This doesn't affect building a jump table, but it affects isDense.)

Hmm... actually, more generally, whether we transform a switch condition is to some extent independent of whether we form a jump table; for example, if we end up with one cluster here, you don't need a jump table at all. Maybe we can leave that for a followup, though.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h

267 ↗

(On Diff #60973)

It's unintuitive to hide this in the constructor... it seems like a good idea to push the Shift == 0 special-case out.

[hans]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1971 ↗

(On Diff #60973)

I was going to comment that the code in the patch summary could use rotation instead, but it seems you're already on it :-)

8434 ↗

(On Diff #60973)

This code doesn't take into account that a range of cases could be transformed to become dense, so I'm not sure your patch will actually find any more jump table lowerings that don't cover the whole switch?

[hans]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

8434 ↗

(On Diff #60973)

I guess what I was saying with this comment is that since you're not going to get the benefit of transforming parts of the switch, it there a point in doing this in the DAG instead of in the IR? I think the rotation trick should work in the IR too..

[eli.friedman]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

8434 ↗

(On Diff #60973)

Err, I guess the rotation trick works in IR? You can let switch lowering do the compare+branch. Not sure what I was thinking before.

That said, we don't really want to do switch lowering in IR. I mean, in theory, we can lower switches in IR, but we want to do all of switch lowering at the same time. Otherwise, you end up with weird guessing games to predict what SelectionDAG will actually do.

[jmolloy]

Hi guys,

Thanks for the comments. I was away on vacation for 2 weeks which is why this has been sitting around.

Hans, if you're OK with it being done in DAG for the reasons mentioned, is this good to commit?

Cheers,

James

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

8434 ↗

(On Diff #60973)

I agree with Eli here. It'd be nice to densify individual switch case ranges, and perhaps we could do that in the future?

The main advantage I got from moving from IR to SDAG for this is the heuristics don't need second-guessing, which was a major improvement. Doing this in IR would give one large advantage - it'd allow the switch table lookup lowering to work too. I intend to implement this directly in the switch table lookup lowering code though, which should be fairly simple.

[jmolloy]

Hi Hans, Sanjoy,

Ping! :) Are you now happy with this?

Cheers,

James

[jmolloy]

Ping!

[hans]

The code seems fine to me. I'm OK with this landing.

I still wonder if we shouldn't just do this in SimplifyCFG though. How big is the pay-off here? I've seen switches with e.g. a common factor of 10 in the cases, but are powers of 2 common enough to justify the complexity, especially if we start looking at densifying only parts of a switch?

I just imagine that doing this in SimplifyCFG instead would be so much simpler: we'd look for a common power-of-two factor and slap a rotate-right operation in front of the switch. This could potentially allow more lookup tables, jump tables, bit-test lowerings, new adjacent cases, who knows -- one could argue that densifying the switch is just general goodness, and if it doesn't pay off it's just an extra rotate which is cheap.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

8389 ↗

(On Diff #60973)

This comment refers to the algorithm that finds dense partitions in Clusters, which is O(n^2). Now that there's more stuff happening below, maybe update it to "don't try any harder at -O0", or you could jsut move your code before it, since I don't think it would be a problem running it a -O0.

[jmolloy]

Hi Hans,

I do agree with you. I think that the major reason I moved from SimplifyCFG to SDAGBuilder was that, because we were inserting new CFG edges, we didn't want to do this unless we could be *sure* that a dense switch lowering was used.

However now we've used the rotate trick, the inserted code is so tiny that we don't have to be so conservative. So I think it's OK to replicate the density heuristic in SimplifyCFG and that's what I've done.

I hope this is now more acceptable to you? :)

Cheers,

James

[hans]

Thanks, I like this :-) Just some nits.

lib/Transforms/Utils/SimplifyCFG.cpp

4974

Aren't Values sorted (from line 5081), so this should be unnecessary?

4997

Maybe we should check DL.fitsInLegalInteger() instead of hard-coding 64 here.

Oh I see, 64 is because we use [u]int64_t below. It might still be a good idea to check fitsInLegalInteger() so we're not creating a non-legal rotate operation.

5018

why isn't Base int64_t already? Also, could the subtraction overflow?

5036

Any bijective transformation of the case values will work. The problem is just deciding when they're profitable, I suppose. But I'm all for more creative switch lowerings :-)

[jmolloy]

Thanks Hans!

New version - hopefully this one should be ready to go?

Cheers,

James

lib/Transforms/Utils/SimplifyCFG.cpp

5018

I don't think overflow is possible. The subtraction moves values towards zero; the only way I can consider it overflowing is if the difference between the smallest value and the largest value is INT64_MAX or greater which is impossible.

5036

:)

[hans]

lgtm

[jmolloy]

This was committed with Hans' LGTM.

lib/Transforms/Utils/SimplifyCFG.cpp

static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   SI->eraseFromParent();
 
   ++NumLookupTables;
   if (NeedMask)
     ++NumLookupTablesHoles;
   return true;
 }
 
+// Try and 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) {
+  if (SI->getCondition()->getType()->getIntegerBitWidth() > 64)
+    return false;
+  // Only bother with this optimization if there are more than 2 switch cases.
+  if (SI->getNumCases() <= 2)
+    return false;
+  SmallVector<uint64_t,4> Values;
+  for (auto &C : SI->cases())
+    Values.push_back(C.getCaseValue()->getValue().getLimitedValue());
+  std::sort(Values.begin(), Values.end());
+
+  bool HasHoles = Values.back() != SI->getNumCases() - 1;
+  if (!HasHoles)
+    // This switch is already dense.
+    return false;
+
+  uint64_t Base = Values[0];
+  for (auto &V : Values)
+    V -= Base;
+  uint64_t Divisor = Values[1];
+  if (Divisor == 1 || !llvm::isPowerOf2_64(Divisor))
+    // We refuse to output a costly urem/udiv pair.
+    return false;
+  
+  unsigned I = 0;
+  for (auto &V : Values)
+    if (V / Divisor != I++)
+      // The switch still has a hole.
+      return false;
+  
+  // We can range-reduce this switch (X) to switch (X - Base) / Divisor.
+  // This should be profitable if the resulting switch has no holes -
+  // this should result in a better lowering.
+  auto *Ty = SI->getCondition()->getType();
+  Builder.SetInsertPoint(SI);
+  auto *Sub = Builder.CreateSub(SI->getCondition(), ConstantInt::get(Ty, Base));
+  auto *Div = Builder.CreateUDiv(Sub, ConstantInt::get(Ty, Divisor));
+  auto *Rem = Builder.CreateURem(Sub, ConstantInt::get(Ty, Divisor));
+  auto *Cmp = Builder.CreateICmpEQ(Rem, ConstantInt::get(Ty, 0));
+  SI->replaceUsesOfWith(SI->getCondition(), Div);
+
+  BasicBlock *OldBB = SI->getParent();
+  SplitBlock(OldBB, SI);
+  OldBB->getTerminator()->eraseFromParent();
+  BranchInst::Create(SI->getParent(), SI->getDefaultDest(), Cmp, OldBB);
+
+  // Correct any PHIs at the destination.
+  for (auto &I : *SI->getDefaultDest()) {
+    auto *PN = dyn_cast<PHINode>(&I);
+    if (!PN)
+      break;
+    auto *V = PN->getIncomingValueForBlock(SI->getParent());
+    PN->addIncoming(V, OldBB);
+  }
+
+  for (auto &C : SI->cases()) {
+    auto *Orig = C.getCaseValue();
+    auto BW = Orig->getValue().getBitWidth();
+    auto Sub = Orig->getValue() - APInt(BW, Base);
+    auto *New = ConstantInt::get(Ty, Sub.udiv(APInt(BW, Divisor)));
+    SI->replaceUsesOfWith(Orig, New);
+  }
+  return true;
+}
+
 bool SimplifyCFGOpt::SimplifySwitch(SwitchInst *SI, IRBuilder<> &Builder) {
   BasicBlock *BB = SI->getParent();
 
   if (isValueEqualityComparison(SI)) {
     // If we only have one predecessor, and if it is a branch on this value,
     // see if that predecessor totally determines the outcome of this switch.
     if (BasicBlock *OnlyPred = BB->getSinglePredecessor())
       if (SimplifyEqualityComparisonWithOnlyPredecessor(SI, OnlyPred, Builder))
         return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
     Value *Cond = SI->getCondition();
     if (SelectInst *Select = dyn_cast<SelectInst>(Cond))
       if (SimplifySwitchOnSelect(SI, Select))
         return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
     // If the block only contains the switch, see if we can fold the block
     // away into any preds.
     BasicBlock::iterator BBI = BB->begin();
     // Ignore dbg intrinsics.
     while (isa<DbgInfoIntrinsic>(BBI))
       ++BBI;
     if (SI == &*BBI)
       if (FoldValueComparisonIntoPredecessors(SI, Builder))
         return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
   }
 
   // Try to transform the switch into an icmp and a branch.
   if (TurnSwitchRangeIntoICmp(SI, Builder))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   // Remove unreachable cases.
   if (EliminateDeadSwitchCases(SI, AC, DL))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   if (SwitchToSelect(SI, Builder, AC, DL))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   if (ForwardSwitchConditionToPHI(SI))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   if (SwitchToLookupTable(SI, Builder, DL, TTI))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
+  if (ReduceSwitchRange(SI, Builder, DL, TTI))
+    return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
+
   return false;
 }
 
 bool SimplifyCFGOpt::SimplifyIndirectBr(IndirectBrInst *IBI) {
   BasicBlock *BB = IBI->getParent();
   bool Changed = false;
 
   // Eliminate redundant destinations.

test/Transforms/SimplifyCFG/rangereduce.ll

+; RUN: opt < %s -simplifycfg -S | FileCheck %s
+
+; CHECK-LABEL: @test1
+; CHECK: %1 = sub i32 %a, 97
+; CHECK: %2 = udiv i32 %1, 2
+; CHECK: %3 = urem i32 %1, 2
+; CHECK: %4 = icmp eq i32 %3, 0
+; CHECK: br i1 %4, label %.split, label %def
+; CHECK:.split:
+; CHECK:  switch i32 %2, label %def [
+; CHECK:    i32 0, label %one
+; CHECK:    i32 1, label %two
+; CHECK:    i32 2, label %three
+; CHECK:  ]
+; CHECK: phi i32 [ 8867, %.split ], [ 11984, %one ], [ 1143, %two ], [ 99783, %three ], [ 8867, %0 ]
+define i32 @test1(i32 %a) {
+  switch i32 %a, label %def [
+    i32 97, label %one
+    i32 99, label %two
+    i32 101, label %three
+  ]
+
+def:
+  ret i32 8867
+
+one:
+  ret i32 11984
+two:
+  ret i32 1143
+three:
+  ret i32 99783
+}
+
+; Optimization shouldn't trigger; bitwidth > 64
+; CHECK-LABEL: @test2
+; CHECK: switch i128 %a, label %def
+define i128 @test2(i128 %a) {
+  switch i128 %a, label %def [
+    i128 97, label %one
+    i128 99, label %two
+    i128 101, label %three
+  ]
+
+def:
+  ret i128 8867
+
+one:
+  ret i128 11984
+two:
+  ret i128 1143
+three:
+  ret i128 99783
+}
+
+
+; Optimization shouldn't trigger; no holes present
+; CHECK-LABEL: @test3
+; CHECK: switch i32 %a, label %def
+define i32 @test3(i32 %a) {
+  switch i32 %a, label %def [
+    i32 97, label %one
+    i32 98, label %two
+    i32 99, label %three
+  ]
+
+def:
+  ret i32 8867
+
+one:
+  ret i32 11984
+two:
+  ret i32 1143
+three:
+  ret i32 99783
+}
+
+; Optimization shouldn't trigger; not an arithmetic progression
+; CHECK-LABEL: @test4
+; CHECK: switch i32 %a, label %def
+define i32 @test4(i32 %a) {
+  switch i32 %a, label %def [
+    i32 97, label %one
+    i32 99, label %two
+    i32 100, label %three
+  ]
+
+def:
+  ret i32 8867
+
+one:
+  ret i32 11984
+two:
+  ret i32 1143
+three:
+  ret i32 99783
+}
+
+; Optimization shouldn't trigger; not a power of two
+; CHECK-LABEL: @test5
+; CHECK: switch i32 %a, label %def
+define i32 @test5(i32 %a) {
+  switch i32 %a, label %def [
+    i32 97, label %one
+    i32 100, label %two
+    i32 103, label %three
+  ]
+
+def:
+  ret i32 8867
+
+one:
+  ret i32 11984
+two:
+  ret i32 1143
+three:
+  ret i32 99783
+}