[Transform] Rewrite LowerSwitch using APInt

[Peter]

Summary

This rewrite fixes https://github.com/llvm/llvm-project/issues/59316.

Previously LowerSwitch uses int64_t, which will crash on case branches using integers with more than 64 bits.
Using APInt fixes this problem. This patch also includes a test

[arsenm]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

277

This was >= before

328–331

const ref?

436

APSint?

llvm/test/Transforms/LowerSwitch/pr59316.ll

1

Needs checks

3

Use opaque pointers

14

Don't use anonymous values

18

Does this test cover all the paths?

[Peter]

1. update tests
2. use APSInt
3. fix some other problems I forgot to fix.

[Peter]

use const ref when possible

[nikic]

Why do some parts of this use APSInt?

[RKSimon]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

358

Please remove / pre-commit any clang-format only changes to untouched code

442

Do we not have a static helper to create values like this?

[Peter]

1. precommit format
2. using APSInt static constructor for SignedMin and SignedMax

[Peter]

In D140747#4019289, @nikic wrote:

Why do some parts of this use APSInt?

The old code assumes the branch case to be signed integer (int64_t), therefore some place use APSInt to make signess more clear.

[nikic]

In D140747#4022123, @Peter wrote:

In D140747#4019289, @nikic wrote:

Why do some parts of this use APSInt?

The old code assumes the branch case to be signed integer (int64_t), therefore some place use APSInt to make signess more clear.

APSInt is only used if you need to dynamically switch between signed and unsigned operations. Otherwise APInt with appropriate operations is used -- there are signed and unsigned variants for all operations where it makes a difference.

[Peter]

use APInt instead of APSInt

[RKSimon]

a couple of minor style comments

llvm/lib/Transforms/Utils/LowerSwitch.cpp

462

(style) assert message

490–491

(style) Assert message

[Peter]

add assert info

[Peter]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

490–491

Done

[RKSimon]

LGTM

[uabelho]

Hello,

The following starts crashing with this patch:

opt -passes=lowerswitch bbi-77637.ll -o /dev/null

It fails with

opt: ../lib/Transforms/Utils/LowerSwitch.cpp:461: void (anonymous namespace)::ProcessSwitchInst(llvm::SwitchInst *, SmallPtrSetImpl<llvm::BasicBlock *> &, llvm::AssumptionCache *, llvm::LazyValueInfo *): Assertion `N.sge(SignedZero) && "Popularity shouldn't be negative."' failed.

bbi-77637.ll261 BDownload

[RKSimon]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

462

It looks like this can't handle i1 types properly

[Peter]

Not just i1 but for any type. (Easiest to trigger on i1 tho).

The root cause is N = High - Low + 1 may overflow to a negative number. (Imagine i32 N = i32::max - i32::min + 1 as signed.)
Promoting N and Pop to higher precision would help. I am working on it.

[Peter]

Popularity is now represented using higher precision unsigned APInt to avoid overflow

[Peter]

add comments

[arsenm]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

461

I don't understand why this wouldn't just be unsigned to begin with

[Peter]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

461

Old version used int64_t N, which would have the same problem.

[uabelho]

We see the

Assertion `N.sge(SignedZero) && "Popularity shouldn't be negative."' failed.

crash fairly often in downstream fuzz testing.
Is someone planning to push a fix?

[Peter]

In D140747#4068376, @uabelho wrote:

We see the

Assertion `N.sge(SignedZero) && "Popularity shouldn't be negative."' failed.

crash fairly often in downstream fuzz testing.
Is someone planning to push a fix?

I think this problems have already been fixed in this patch. I am waiting for someone to review it before I patch it.

[arsenm]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

461

Right, so why isn’t it OK to just switch to always using unsigned?

[Peter]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

461

Well, old version has existed and works for years, I figure we shouldn't change that decision unless we have a good reason

And I now this (fixed) bug gave us a reason. :)

[RKSimon]

LGTM

[uabelho]

Thanks!

llvm/lib/Transforms/Utils/LowerSwitch.cpp

BasicBlock *SwitchConvert(CaseItr Begin, CaseItr End, ConstantInt *LowerBound,
   ConstantInt *NewUpperBound = ConstantInt::get(NewLowerBound->getContext(),
                                                 NewLowerBound->getValue() - 1);
 
   if (!UnreachableRanges.empty()) {
     // Check if the gap between LHS's highest and NewLowerBound is unreachable.
     APInt GapLow = LHS.back().High->getValue() + 1;
     APInt GapHigh = NewLowerBound->getValue() - 1;
     IntRange Gap = {GapLow, GapHigh};
     if (GapHigh.sge(GapLow) && IsInRanges(Gap, UnreachableRanges))
       NewUpperBound = LHS.back().High;
   }
 
   LLVM_DEBUG(dbgs() << "LHS Bounds ==> [" << LowerBound->getValue() << ", "
                     << NewUpperBound->getValue() << "]\n"
                     << "RHS Bounds ==> [" << NewLowerBound->getValue() << ", "
                     << UpperBound->getValue() << "]\n");
 
unsigned Clusterify(CaseVector &Cases, SwitchInst *SI) {
   }
 
   llvm::sort(Cases, CaseCmp());
 
   // Merge case into clusters
   if (Cases.size() >= 2) {
     CaseItr I = Cases.begin();
     for (CaseItr J = std::next(I), E = Cases.end(); J != E; ++J) {
       const APInt &nextValue = J->Low->getValue();
       const APInt &currentValue = I->High->getValue();
       BasicBlock *nextBB = J->BB;
       BasicBlock *currentBB = I->BB;
 
       // If the two neighboring cases go to the same destination, merge them
       // into a single case.
       assert(nextValue.sgt(currentValue) &&
              "Cases should be strictly ascending");
       if ((nextValue == currentValue + 1) && (currentBB == nextBB)) {
         I->High = J->High;
         // FIXME: Combine branch weights.

 /// Replace the specified switch instruction with a sequence of chained if-then
 /// insts in a balanced binary search.
 void ProcessSwitchInst(SwitchInst *SI,
                        SmallPtrSetImpl<BasicBlock *> &DeleteList,
                        AssumptionCache *AC, LazyValueInfo *LVI) {
   BasicBlock *OrigBlock = SI->getParent();
   Function *F = OrigBlock->getParent();
   Value *Val = SI->getCondition(); // The value we are switching on...
   BasicBlock *Default = SI->getDefaultDest();
 
   // Don't handle unreachable blocks. If there are successors with phis, this
   // would leave them behind with missing predecessors.
   if ((OrigBlock != &F->getEntryBlock() && pred_empty(OrigBlock)) ||
       OrigBlock->getSinglePredecessor() == OrigBlock) {
     DeleteList.insert(OrigBlock);
     return;
   }
 
   // Prepare cases vector.
   CaseVector Cases;
   const unsigned NumSimpleCases = Clusterify(Cases, SI);
   IntegerType *IT = cast<IntegerType>(SI->getCondition()->getType());
   const unsigned BitWidth = IT->getBitWidth();
-  APInt SignedZero(BitWidth, 0);
+  // Explictly use higher precision to prevent unsigned overflow where
+  // `UnsignedMax - 0 + 1 == 0`
+  APInt UnsignedZero(BitWidth + 1, 0);
   APInt UnsignedMax = APInt::getMaxValue(BitWidth);
   LLVM_DEBUG(dbgs() << "Clusterify finished. Total clusters: " << Cases.size()
                     << ". Total non-default cases: " << NumSimpleCases
                     << "\nCase clusters: " << Cases << "\n");
 
   // If there is only the default destination, just branch.
   if (Cases.empty()) {
     BranchInst::Create(Default, OrigBlock);
if (isa<UnreachableInst>(Default->getFirstNonPHIOrDbg())) {
     UpperBound = ConstantInt::get(SI->getContext(), Max);
     DefaultIsUnreachableFromSwitch = (Min + (NumSimpleCases - 1) == Max);
   }
 
   std::vector<IntRange> UnreachableRanges;
 
   if (DefaultIsUnreachableFromSwitch) {
     DenseMap<BasicBlock *, APInt> Popularity;
-    APInt MaxPop(SignedZero);
+    APInt MaxPop(UnsignedZero);
     BasicBlock *PopSucc = nullptr;
 
     APInt SignedMax = APInt::getSignedMaxValue(BitWidth);
     APInt SignedMin = APInt::getSignedMinValue(BitWidth);
     IntRange R = {SignedMin, SignedMax};
     UnreachableRanges.push_back(R);
     for (const auto &I : Cases) {
       const APInt &Low = I.Low->getValue();
       const APInt &High = I.High->getValue();
 
       IntRange &LastRange = UnreachableRanges.back();
       if (LastRange.Low.eq(Low)) {
         // There is nothing left of the previous range.
         UnreachableRanges.pop_back();
       } else {
         // Terminate the previous range.
         assert(Low.sgt(LastRange.Low));
         LastRange.High = Low - 1;
       }
       if (High.ne(SignedMax)) {
         IntRange R = {High + 1, SignedMax};
         UnreachableRanges.push_back(R);
       }
 
       // Count popularity.
-      APInt N = High - Low + 1;
-      assert(N.sge(SignedZero) && "Popularity shouldn't be negative.");
+      assert(High.sge(Low) && "Popularity shouldn't be negative.");
+      APInt N = High.sext(BitWidth + 1) - Low.sext(BitWidth + 1) + 1;
       // Explict insert to make sure the bitwidth of APInts match
-      APInt &Pop = Popularity.insert({I.BB, APInt(SignedZero)}).first->second;
-      if ((Pop += N).sgt(MaxPop)) {
+      APInt &Pop = Popularity.insert({I.BB, APInt(UnsignedZero)}).first->second;
+      if ((Pop += N).ugt(MaxPop)) {
         MaxPop = Pop;
         PopSucc = I.BB;
       }
     }
 #ifndef NDEBUG
     /* UnreachableRanges should be sorted and the ranges non-adjacent. */
     for (auto I = UnreachableRanges.begin(), E = UnreachableRanges.end();
          I != E; ++I) {
       assert(I->Low.sle(I->High));
       auto Next = I + 1;
       if (Next != E) {
         assert(Next->Low.sgt(I->High));
       }
     }
 #endif
 
     // As the default block in the switch is unreachable, update the PHI nodes
     // (remove all of the references to the default block) to reflect this.
     const unsigned NumDefaultEdges = SI->getNumCases() + 1 - NumSimpleCases;
     for (unsigned I = 0; I < NumDefaultEdges; ++I)
       Default->removePredecessor(OrigBlock);
 
     // Use the most popular block as the new default, reducing the number of
     // cases.
-    assert(MaxPop.sgt(SignedZero) && PopSucc &&
-           "Max populartion shouldn't be negative.");
     Default = PopSucc;
     llvm::erase_if(Cases,
                    [PopSucc](const CaseRange &R) { return R.BB == PopSucc; });
 
     // If there are no cases left, just branch.
     if (Cases.empty()) {
       BranchInst::Create(Default, OrigBlock);
       SI->eraseFromParent();
       // As all the cases have been replaced with a single branch, only keep
       // one entry in the PHI nodes.
-      for (APInt I(SignedZero); I.slt(MaxPop - 1); ++I)
+      if (!MaxPop.isZero())
+        for (APInt I(UnsignedZero); I.ult(MaxPop - 1); ++I)
           PopSucc->removePredecessor(OrigBlock);
       return;
     }
 
     // If the condition was a PHI node with the switch block as a predecessor
     // removing predecessors may have caused the condition to be erased.
     // Getting the condition value again here protects against that.
     Val = SI->getCondition();
   }

llvm/test/Transforms/LowerSwitch/pr59316.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=lowerswitch -S | FileCheck %s
 
 define i64 @f(i1 %bool, i128 %i128) {
 ; CHECK-LABEL: @f(
 ; CHECK-NEXT:  BB:
 ; CHECK-NEXT:    br label [[NODEBLOCK1:%.*]]
 ; CHECK:       NodeBlock1:
 ; CHECK-NEXT:    [[PIVOT2:%.*]] = icmp slt i128 [[I128:%.*]], 16201310291018008446
 ; CHECK-NEXT:    br i1 [[PIVOT2]], label [[LEAFBLOCK:%.*]], label [[NODEBLOCK:%.*]]
 ; CHECK:       NodeBlock:
 ; CHECK-NEXT:    [[PIVOT:%.*]] = icmp slt i128 [[I128]], 16201310291018008447
 ; CHECK-NEXT:    br i1 [[PIVOT]], label [[SW_C3:%.*]], label [[SW_C2:%.*]]
 ; CHECK:       LeafBlock:
 ; CHECK-NEXT:    [[SWITCHLEAF:%.*]] = icmp eq i128 [[I128]], 16201310291018008445
 ; CHECK-NEXT:    br i1 [[SWITCHLEAF]], label [[SW_C4:%.*]], label [[SW_C1:%.*]]
 ; CHECK:       BB1:
 ; CHECK-NEXT:    unreachable
 ; CHECK:       SW_C1:
 ; CHECK-NEXT:    br i1 [[BOOL:%.*]], label [[BB1:%.*]], label [[SW_C1]]
 ; CHECK:       SW_C2:
 ; CHECK-NEXT:    ret i64 0
 ; CHECK:       SW_C3:
 ; CHECK-NEXT:    ret i64 1
 ; CHECK:       SW_C4:
 ; CHECK-NEXT:    ret i64 2

 ; CHECK-NEXT:    unreachable
 ;
 BB:
   switch i128 %i128, label %BB1 []
 
 BB1:                                              ; preds = %BB
   unreachable
 }
+
+define void @f_i1() {
+entry:
+  switch i1 false, label %sw.bb [
+    i1 false, label %sw.bb12
+  ]
+
+sw.bb:                                            ; preds = %entry
+  unreachable
+
+sw.bb12:                                          ; preds = %entry
+  unreachable
+}
+
+define void @f_i2(i2 %cond) {
+entry:
+  switch i2 %cond, label %sw.bb [
+    i2 0, label %sw.bb12
+    i2 1, label %sw.bb12
+    i2 2, label %sw.bb12
+    i2 3, label %sw.bb12
+  ]
+
+sw.bb:                                            ; preds = %entry
+  unreachable
+
+sw.bb12:                                          ; preds = %entry
+  unreachable
+}
 No newline at end of file