[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

279

This was >= before

330

const ref?

432

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

360

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

440

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

466

(style) assert message

494

(style) Assert message

[Peter]

add assert info

[Peter]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

494

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

466

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

456–457

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

[Peter]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

456–457

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

456–457

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

[Peter]

llvm/lib/Transforms/Utils/LowerSwitch.cpp

456–457

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

 #include <vector>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "lower-switch"
 
 namespace {
 
-  struct IntRange {
-    int64_t Low, High;
-  };
+  APInt Low, High;
 
 } // end anonymous namespace
 
 namespace {
 // Return true iff R is covered by Ranges.
 bool IsInRanges(const IntRange &R, const std::vector<IntRange> &Ranges) {
   // Note: Ranges must be sorted, non-overlapping and non-adjacent.
 
   // Find the first range whose High field is >= R.High,
   // then check if the Low field is <= R.Low. If so, we
   // have a Range that covers R.
   auto I = llvm::lower_bound(
-      Ranges, R, [](IntRange A, IntRange B) { return A.High < B.High; });
-  return I != Ranges.end() && I->Low <= R.Low;
+      Ranges, R, [](IntRange A, IntRange B) { return A.High.slt(B.High); });
+  return I != Ranges.end() && I->Low.sle(R.Low);
 }
 
 struct CaseRange {
   ConstantInt *Low;
   ConstantInt *High;
   BasicBlock *BB;
 
   CaseRange(ConstantInt *low, ConstantInt *high, BasicBlock *bb)
BasicBlock *SwitchConvert(CaseItr Begin, CaseItr End, ConstantInt *LowerBound,
 
   // Because NewLowerBound is never the smallest representable integer
   // it is safe here to subtract one.
   ConstantInt *NewUpperBound = ConstantInt::get(NewLowerBound->getContext(),
                                                 NewLowerBound->getValue() - 1);
 
   if (!UnreachableRanges.empty()) {
     // Check if the gap between LHS's highest and NewLowerBound is unreachable.
-    int64_t GapLow = LHS.back().High->getSExtValue() + 1;
-    int64_t GapHigh = NewLowerBound->getSExtValue() - 1;
-    IntRange Gap = { GapLow, GapHigh };
-    if (GapHigh >= GapLow && IsInRanges(Gap, UnreachableRanges))
+    APInt GapLow = LHS.back().High->getValue() + 1;
+    APInt GapHigh = NewLowerBound->getValue() - 1;
+    IntRange Gap = {GapLow, GapHigh};
+    if (GapHigh.sgt(GapLow) && IsInRanges(Gap, UnreachableRanges))
       NewUpperBound = LHS.back().High;
   }
 
   LLVM_DEBUG(dbgs() << "LHS Bounds ==> [" << LowerBound->getSExtValue() << ", "
                     << NewUpperBound->getSExtValue() << "]\n"
                     << "RHS Bounds ==> [" << NewLowerBound->getSExtValue()
                     << ", " << UpperBound->getSExtValue() << "]\n");
 
   // Create a new node that checks if the value is < pivot. Go to the
   // left branch if it is and right branch if not.
-  Function* F = OrigBlock->getParent();
-  BasicBlock* NewNode = BasicBlock::Create(Val->getContext(), "NodeBlock");
+  Function *F = OrigBlock->getParent();
+  BasicBlock *NewNode = BasicBlock::Create(Val->getContext(), "NodeBlock");
 
-  ICmpInst* Comp = new ICmpInst(ICmpInst::ICMP_SLT,
-                                Val, Pivot.Low, "Pivot");
+  ICmpInst *Comp = new ICmpInst(ICmpInst::ICMP_SLT, Val, Pivot.Low, "Pivot");
 
   BasicBlock *LBranch =
       SwitchConvert(LHS.begin(), LHS.end(), LowerBound, NewUpperBound, Val,
                     NewNode, OrigBlock, Default, UnreachableRanges);
   BasicBlock *RBranch =
       SwitchConvert(RHS.begin(), RHS.end(), NewLowerBound, UpperBound, Val,
                     NewNode, OrigBlock, Default, UnreachableRanges);
 
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) {
-      int64_t nextValue = J->Low->getSExtValue();
-      int64_t currentValue = I->High->getSExtValue();
-      BasicBlock* nextBB = J->BB;
-      BasicBlock* currentBB = I->BB;
+      APInt nextValue = J->Low->getValue();
+      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 > currentValue && "Cases should be strictly ascending");
+      assert(nextValue.sgt(currentValue) &&
+             "Cases should be strictly ascending");
       if ((nextValue == currentValue + 1) && (currentBB == nextBB)) {
         I->High = J->High;
         // FIXME: Combine branch weights.
       } else if (++I != J) {
         *I = *J;
       }
     }
     Cases.erase(std::next(I), Cases.end());
   }
 
   return NumSimpleCases;
 }
 
 /// 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();
+  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;
   }
if (isa<UnreachableInst>(Default->getFirstNonPHIOrDbg())) {
     LowerBound = ConstantInt::get(SI->getContext(), Min);
     UpperBound = ConstantInt::get(SI->getContext(), Max);
     DefaultIsUnreachableFromSwitch = (Min + (NumSimpleCases - 1) == Max);
   }
 
   std::vector<IntRange> UnreachableRanges;
 
   if (DefaultIsUnreachableFromSwitch) {
-    DenseMap<BasicBlock *, unsigned> Popularity;
-    unsigned MaxPop = 0;
+    unsigned BitWidth = Cases[0].Low->getBitWidth();
+    APInt SignedZero(BitWidth, 0, true);
+    DenseMap<BasicBlock *, APInt> Popularity;
+    APInt MaxPop(SignedZero);
     BasicBlock *PopSucc = nullptr;
 
-    IntRange R = {std::numeric_limits<int64_t>::min(),
-                  std::numeric_limits<int64_t>::max()};
+    APInt MaxCase(SignedZero);
+    // Signed max: 0b01111....
+    MaxCase.setAllBits();
+    MaxCase.setBitVal(BitWidth - 1, 0);
+    // Signed min: 0b10000....
+    APInt MinCase(SignedZero);
+    MinCase.setBitVal(BitWidth - 1, 1);
+    IntRange R = {MinCase, MaxCase};
     UnreachableRanges.push_back(R);
     for (const auto &I : Cases) {
-      int64_t Low = I.Low->getSExtValue();
-      int64_t High = I.High->getSExtValue();
+      APInt Low = I.Low->getValue();
+      APInt High = I.High->getValue();
 
       IntRange &LastRange = UnreachableRanges.back();
-      if (LastRange.Low == Low) {
+      if (LastRange.Low.eq(Low)) {
         // There is nothing left of the previous range.
         UnreachableRanges.pop_back();
       } else {
         // Terminate the previous range.
-        assert(Low > LastRange.Low);
+        assert(Low.sgt(LastRange.Low));
         LastRange.High = Low - 1;
       }
-      if (High != std::numeric_limits<int64_t>::max()) {
-        IntRange R = { High + 1, std::numeric_limits<int64_t>::max() };
+      if (High.ne(MaxCase)) {
+        IntRange R = {High + 1, MaxCase};
         UnreachableRanges.push_back(R);
       }
 
       // Count popularity.
-      int64_t N = High - Low + 1;
-      unsigned &Pop = Popularity[I.BB];
-      if ((Pop += N) > MaxPop) {
+      APInt N = High - Low + 1;
+      assert(N.sge(SignedZero));
+      // 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)) {
         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 <= I->High);
+      assert(I->Low.sle(I->High));
       auto Next = I + 1;
       if (Next != E) {
-        assert(Next->Low > I->High);
+        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 > 0 && PopSucc);
+    assert(MaxPop.sgt(SignedZero) && PopSucc);
     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 (unsigned I = 0 ; I < (MaxPop - 1) ; ++I)
+      for (APInt I(SignedZero); I.slt(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

+; RUN: opt < %s -passes=lowerswitch
+
+define i64 @f(i1 %0, i128 %1, i64* %2) {
+BB:
+  switch i128 %1, label %BB1 [
+    i128 627, label %SW_C
+    i128 16201310291018008447, label %SW_C2
+  ]
+
+BB1:                                              ; preds = %SW_C, %BB
+  unreachable
+
+SW_C:                                             ; preds = %SW_C, %BB
+  br i1 %0, label %BB1, label %SW_C
+
+SW_C2:                                            ; preds = %BB
+  ret i64 0
+}