Switch optimization in IR for known maximum switch values

[ayonam]

Summary

This patch is for optimizing a switch statement where the switch expression can be analyzed to determine the maximum value it can take. In such cases, the default statement generates a conditional branch comparing the switch value with the highest value of case statement given. This branch can be avoided by adding the missing case statements between the highest value of case statement given and the maximum value the switch expression can take and then redirecting all these new case statements to the default block. Thereafter, the default block can be marked as unreachable, which would remove the conditional branch while lowering the switch statement. This part is dealt with in the Revision D52002 that is also up for review.

[ayonam]

Hans,

What do those yellow lines on the side of the diff denote? Am I supposed to fix something there? Does that denote some kind of non-compliance to the coding standards that were detected?

[efriedma]

The yellow bars are just Phabricator noting copy-pasted code; if you hover over the bar, you can get more details from the tooltip. You might not have to do anything about it in general, but in this case I think it's worth refactoring.

lib/Transforms/Utils/SimplifyCFG.cpp

4866

Can this shift overflow?

[ayonam]

In D52707#1251513, @efriedma wrote:

The yellow bars are just Phabricator noting copy-pasted code; if you hover over the bar, you can get more details from the tooltip. You might not have to do anything about it in general, but in this case I think it's worth refactoring.

Thanks. I will refactor the second block of code since the first block has one additional operation of storing the PHIs in a vector.

As for the other comment in the code:

countMaxTrailingOnes will not exceed the number of bits in the switch expression type (local variable ValueWidth above). But then the '1' should probably be '1LL' and MaxSwitchValue should be of type 'long long' to ensure that all possible value widths are taken care of.

lib/Transforms/Utils/SimplifyCFG.cpp

4866

countMaxTrailingOnes will not exceed the number of bits in the switch expression type (local variable ValueWidth above). But then the '1' should probably be '1LL' and MaxSwitchValue should be of type 'long long' to ensure that all possible value widths are taken care of.

[efriedma]

lib/Transforms/Utils/SimplifyCFG.cpp

4866

The LLVM IR "switch" instruction allows the operand to be any width. (And even in C, you can write a switch using __uint128_t.) You probably want to use APInt here.

And actually, I'm not sure why you're shifting here in the first place; the maximum possible value is exactly "~Known.Zero".

[ayonam]

lib/Transforms/Utils/SimplifyCFG.cpp

4866

Yes, I didn't factor in the __uint128_t type. I will change the code to use the ~Known.Zero value for the Max Switch Value.

Using the APInt type would be safer. I will make the necessary changes.

[ayonam]

I have implemented the review comments given earlier. The updated patch is attached herewith. Kindly review the same.

Thanks

[lebedev.ri]

No tests?

[hans]

Thanks for doing this! I have some comments, and as was pointed out, this needs tests.

lib/Transforms/Utils/SimplifyCFG.cpp

4851

The way I think about this is that it "widens" the switch to cover all input values. So maybe call the function something like that?

4866

I think the following code would be easier to follow if you do these "known value is not negative etc." checks and then return false if they don't pass. Same thing for the unreachable default check: if that's true, there's no point in continuing.

4877

Instead of using floating point for this, how about something like "NumCases * 100 / MaxSwitchValue >= 70"?

4880

I think we need to care about the holes too, because we're going to mark the current default unreachable..

4891

What's happening here? Why do we need to find a constant value on the phi node?

We'll need to update any phi node in the default block when adding the new case, but there's no need for it to be constant.

4908

I don't think this is enough. I think we need to fill the "holes" too.

4920

This needs a comment explaining that the new default is unreachable because the switch now has cases for all possible inputs.

4928

Let's add a statistic variable to count how often this transformation happens (see e.g. NumLinearMaps).

[ayonam]

lib/Transforms/Utils/SimplifyCFG.cpp

4891

It should just be a Value and not a Constant. I had a test case that had a constant on the PHI node and hence I had coded thus.

I tried changing this to a Value and realized that there isn't a method that gives me all the Values that are created in a given Case block. There is a method that returns all constants that are created in a Case block. Would you know of any other method that would return me the Values? Or else I will have to write one of my own.

[ayonam]

I have updated the patch with the changes mentioned by Hans. Please review the same.

[ayonam]

The earlier patch had a block of code that was commented out. This patch removes that block of code. All other changes asked for in the earlier review comments are there. Kindly review the same.

Thanks

[hans]

lib/Transforms/Utils/SimplifyCFG.cpp

4891

What I'm saying is there should be no need to compute the resulting value at phi nodes at all. That stuff is used when replacing switches with a table of constant values, it does not apply here.

[ayonam]

lib/Transforms/Utils/SimplifyCFG.cpp

4891

We need to get all the PHI nodes resulting from this switch and add an incoming edge to all those PHI nodes for every case that we add to this switch. Else the block containing the PHIs will have more predecessors than the number of incoming edges to the PHIs. I'm trying to do essentially that. Could you please suggest a better way to achieve this?

[hans]

lib/Transforms/Utils/SimplifyCFG.cpp

4891

Yes, the phi nodes in DefaultBlock need to be updated, but this is not the way to do it. I think something like

AddPredecessorToBlock(DefaultBlock, SI->getParent(), SI->getParent())

when adding the cases below should work.

[ayonam]

lib/Transforms/Utils/SimplifyCFG.cpp

4891

Thanks Hans for this tip. I'll update the code and resubmit the patch.

[ayonam]

@hans

The earlier code had a problem in that if we converted every switch that had an unreachable default, to omit the check for the default values beyond the maximum switch case value given in the code, then there were some corner cases where the code generated was semantically incorrect and resulted in runtime failures. These failures are happening while compiling LLVM code itself.

So I have added a boolean in the SwitchInst class to mark switches that have been widened, so that we omit the branch only for those switches. I need to explore the reasons for the behaviour that was observed when we were blindly removing the branch for all switches that had an unreachable default and will post a fix in a later patch.

Please review this and the other two patches 52002 and another one with the test case for this patch.

[lebedev.ri]

In D52707#1332374, @ayonam wrote:

...

Please review this and the other two patches 52002 and another one with the test case for this patch.

Both the D52002 and this one, still do not have any tests..

[ayonam]

In D52707#1332376, @lebedev.ri wrote:

In D52707#1332374, @ayonam wrote:

...

Please review this and the other two patches 52002 and another one with the test case for this patch.

Both the D52002 and this one, still do not have any tests..

I have added the test case for review under D55742. Please review the same. Thanks.

[ayonam]

Reversed the recording of the switch widening in the IR. Now we rely solely on the default being made unreachable.

[ayonam]

Moved the test case from D55742 to this patch. Simplified the test case.

[ayonam]

A gentle reminder to the reviewers for review comments on this please.

Thanks.

[ayonam]

A gentle reminder to the reviewers for review comments on this please. If there are no concerns, then can we move this to the accepted state please?

Thanks.

[ayonam]

@hans

Could you please review this patch after the changes I made? If things are fine, then can this be accepted for commit?

[ayonam]

ping

@hans

Sorry to bother you again. Will it be possible to review this change and accept it please?

Thanks and regards
Ayonam

[hans]

Sorry for the slow response time. I'm back from vacation now.

lib/Transforms/Utils/SimplifyCFG.cpp

4850

The function description is still a bit confusing. I understand what you mean by "remap the known default values into individual case statements", but requires some thinking to understand.

I would suggest rephrasing it as "Possibly widen the switch by adding more cases so that all possible input values are covered, making the default case unreachable."

4860

Since computeKnownBits takes some work to do, I would move the check and return for unreachable default right to the start of the function so we don't do any unnecessary work in that case.

4862

Why do we only care about non-negative values?

And more importantly, why the check for minimum number of leading zeros and maximum trailing ones... just reading that, it sounds like it's always true? Is there any situation when it's not true?

4883

Why does MaxSwitchValue have to be greater than the number of cases?

I'm not sure I'm following the logic here actually. Also, could we have cases that are greater than MaxSwitchValue, or will some other pass have removed them?

lib/Transforms/Utils/SimplifyCFG.cpp

 STATISTIC(NumLinearMaps,
           "Number of switch instructions turned into linear mapping");
 STATISTIC(NumLookupTables,
           "Number of switch instructions turned into lookup tables");
 STATISTIC(
     NumLookupTablesHoles,
     "Number of switch instructions turned into lookup tables (holes checked)");
 STATISTIC(NumTableCmpReuses, "Number of reused switch table lookup compares");
+STATISTIC(NumSwitchWidened,
+          "Number of switch instructions widened to cover all possible values");
 STATISTIC(NumSinkCommons,
           "Number of common instructions sunk down to the end block");
 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
for (unsigned i = 0, e = SI->getNumSuccessors(); i < e; ++i) {
 
     if (Succ == PHI->getParent())
       continue;
     Succ->removePredecessor(SelectBB);
   }
   SI->eraseFromParent();
 }
 
+// Analyze the switch expression and compute the highest possible value of the
+// expression and if known then remap the known default values into individual
+// case statements and direct them to the default block.  Then mark the default
+// destination as unreachable.
+static bool widenSwitchRangeToCoverAllValues(SwitchInst *SI,
+                                             const DataLayout &DL,
+                                             const TargetTransformInfo &TTI) {
+  unsigned ValueWidth = DL.getTypeSizeInBits(SI->getOperand(0)->getType());
+  unsigned OrigValueWidth = ValueWidth;
+  KnownBits Known(ValueWidth);
+  llvm::computeKnownBits(SI->getOperand(0), Known, DL);
+
+  bool UnreachableDefault =
+      isa<UnreachableInst>(SI->getDefaultDest()->getFirstNonPHIOrDbg());
+
+  // We only care for non-negative values.  The number of minimum leading
+  // zeros plus the number of maximum trailing ones should be equal to the
+  // width of the value.  If the default is unreachable then we skip this
+  // optimization.
+  if (Known.isNegative()  ||  UnreachableDefault  ||
+      (Known.countMinLeadingZeros() + Known.countMaxTrailingOnes()) !=
+         ValueWidth)
+    return false;
+
+  // ValueWidth should be the minimum width needed to represent
+  // MaxSwitchValue * 100.  Since 100 needs 7 bits we add 7 to ValueWidth.
+  ValueWidth += 7;
+
+  APInt NumCases(ValueWidth, SI->getNumCases());
+  APInt MaxSwitchValue = (~Known.Zero).zext(ValueWidth);
+  APInt Hundred = APInt(ValueWidth, 100);
+  APInt Seventy = APInt(ValueWidth, 70);
+
+  // At least 70% of the possible values must be defined as case statements
+  // so that there is no penalty for those which have few defined case
+  // statements.
+  if (MaxSwitchValue.ugt(NumCases)  &&
+      SI->getNumCases() > 0  &&
+      (NumCases * Hundred).udiv(MaxSwitchValue).uge(Seventy)) {
+    // Here we not only need to care for those values that are beyond the
+    // highest case value and up to the MaxSwitchValue but also for the other
+    // missing values (holes) that appear between the already defined cases.
+    BasicBlock *DefaultBlock = SI->getDefaultDest();
+    assert(SI->case_begin() != SI->case_end());
+    SmallVector<ConstantInt *, 16> Cases;
+
+    for (auto Case : SI->cases())
+      Cases.push_back(Case.getCaseValue());
+
+    // The function CasesAreContiguous sorts the vector Cases in descending
+    // order.
+    if (CasesAreContiguous(Cases)) {
+      for (APInt I = (Cases[0]->getValue().zext(ValueWidth) + 1);
+           I.ule(MaxSwitchValue);  I++) {
+        SI->addCase(ConstantInt::get(DefaultBlock->getContext(),
+                    I.trunc(OrigValueWidth)), DefaultBlock);
+        if (MaxSwitchValue.ne(I)) {
+            AddPredecessorToBlock(DefaultBlock, SI->getParent(), SI->getParent());
+        }
+      }
+    } else {
+      unsigned int i = 0;
+      bool SeenFirstDefaultCase = false;
+      for (APInt I = MaxSwitchValue, E(ValueWidth, -1);  I.ne(E);  I--) {
+        if (i >= Cases.size()  ||
+            Cases[i]->getValue().zext(ValueWidth) != I) {
+          SI->addCase(ConstantInt::get(DefaultBlock->getContext(),
+                      I.trunc(OrigValueWidth)), DefaultBlock);
+          if (SeenFirstDefaultCase)
+            AddPredecessorToBlock(DefaultBlock, SI->getParent(), SI->getParent());
+          else
+            SeenFirstDefaultCase = true;
+        } else {
+          i++;
+        }
+      }
+    }
+
+    // We have already added cases for all the possible values including those
+    // that would have mapped to the default case and redirected the newly
+    // added cases to the original default block.  Hence, the new default block 
+    // can be marked as unreachable.
+    BasicBlock *NewDefaultBlock;
+    NewDefaultBlock = BasicBlock::Create(DefaultBlock->getContext(),
+                                         "switch.unreachable.default",
+                                         DefaultBlock->getParent());
+    UnreachableInst *dummyInst LLVM_ATTRIBUTE_UNUSED;
+    dummyInst = new UnreachableInst(NewDefaultBlock->getContext(),
+                                    NewDefaultBlock);
+    SI->setDefaultDest(NewDefaultBlock);
+    ++NumSwitchWidened;
+    return true;
+  }
+
+  return false;
+}
+
 /// 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,
                            const TargetTransformInfo &TTI) {
   Value *const Cond = SI->getCondition();
   PHINode *PHI = nullptr;
if (isValueEqualityComparison(SI)) {
 
     // If the block only contains the switch, see if we can fold the block
     // away into any preds.
     if (SI == &*BB->instructionsWithoutDebug().begin())
       if (FoldValueComparisonIntoPredecessors(SI, Builder))
         return requestResimplify();
   }
 
+  // Try to convert small number of default values to individual case
+  // statements pointing to the default block.
+  if (widenSwitchRangeToCoverAllValues(SI, DL, TTI))
+    return requestResimplify();
+
   // Try to transform the switch into an icmp and a branch.
   if (TurnSwitchRangeIntoICmp(SI, Builder))
     return requestResimplify();
 
   // Remove unreachable cases.
   if (eliminateDeadSwitchCases(SI, Options.AC, DL))
     return requestResimplify();
 

test/Transforms/SimplifyCFG/switch-widen.ll

+; RUN: opt %s -S -simplifycfg | FileCheck %s
+@x = external global i32, align 4
+@y = external global i32, align 4
+@z = external global i32, align 4
+
+define i32 @foo(i4) {
+  switch i4 %0, label %default.block [
+  i4 0, label %case_0
+  i4 -4, label %case_12
+  i4 2, label %case_2
+  i4 3, label %case_3
+  i4 -2, label %case_14
+  i4 -1, label %case_15
+  i4 6, label %case_6
+  i4 7, label %case_7
+  i4 -8, label %case_8
+  i4 -7, label %case_9
+  i4 -6, label %case_10
+  i4 -5, label %case_11
+  ]
+
+case_0:                                           ; preds = %1
+  store i32 425, i32* @x
+  br label %return.block
+
+case_12:                                          ; preds = %1
+  store i32 457, i32* @y
+  br label %return.block
+
+case_2:                                           ; preds = %1
+  store i32 65456, i32* @z
+  br label %return.block
+
+case_3:                                           ; preds = %1
+  store i32 57, i32* @y
+  store i32 25, i32* @x
+  br label %return.block
+
+case_14:                                          ; preds = %1
+  store i32 5456, i32* @z
+  store i32 45, i32* @x
+  br label %return.block
+
+case_15:                                          ; preds = %1
+  store i32 42, i32* @x
+  br label %return.block
+
+case_6:                                           ; preds = %1
+  store i32 47, i32* @y
+  br label %return.block
+
+case_7:                                          ; preds = %1
+  store i32 6556, i32* @z
+  br label %return.block
+
+case_8:                                           ; preds = %1
+  store i32 1425, i32* @x
+  store i32 6546, i32* @z
+  br label %return.block
+
+case_9:                                           ; preds = %1
+  store i32 1457, i32* @y
+  store i32 6545, i32* @z
+  br label %return.block
+
+case_10:                                          ; preds = %1
+  store i32 6456, i32* @z
+  store i32 2457, i32* @y
+  br label %return.block
+
+case_11:                                          ; preds = %1
+  store i32 5425, i32* @x
+  br label %return.block
+
+default.block:                                    ; preds = %1
+  store i32 3409, i32* @z
+  store i32 3409, i32* @x
+  br label %return.block
+
+return.block:                                     ; preds = %default.block, %case_11, %case_10, %case_9, %case_8, %case_7, %case_6, %case_15, %case_14, %case_3, %case_2, %case_12, %case_0
+  %2 = phi i32 [ 567, %default.block ], [ 54325, %case_11 ], [ 44325, %case_10 ], [ 34325, %case_9 ], [ 24325, %case_8 ], [ 14325, %case_7 ], [ 4, %case_6 ], [ 43, %case_15 ], [ 432, %case_14 ], [ 5, %case_3 ], [ 25, %case_2 ], [ 325, %case_12 ], [ 4325, %case_0 ]
+  ret i32 %2
+}
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:    switch i4 %0, label %switch.unreachable.default
+; CHECK:         i4 -3, label %default.block
+; CHECK:         i4 5, label %default.block
+; CHECK:         i4 4, label %default.block
+; CHECK:         i4 1, label %default.block
+; CHECK:         default.block: ; preds = %1, %1, %1, %1
+; CHECK:         switch.unreachable.default: ; preds = %1
+; CHECK-NEXT:    unreachable