[PATCH] D140747: [Transform] Rewrite LowerSwitch using APInt

[Peter Rong via Phabricator via llvm-commits]

Peter updated this revision to Diff 488721.
Peter added a comment.

add comments


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D140747/new/

https://reviews.llvm.org/D140747

Files:
  llvm/lib/Transforms/Utils/LowerSwitch.cpp
  llvm/test/Transforms/LowerSwitch/pr59316.ll


Index: llvm/test/Transforms/LowerSwitch/pr59316.ll
===================================================================
--- llvm/test/Transforms/LowerSwitch/pr59316.ll
+++ llvm/test/Transforms/LowerSwitch/pr59316.ll
@@ -62,3 +62,32 @@
 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
Index: llvm/lib/Transforms/Utils/LowerSwitch.cpp
===================================================================
--- llvm/lib/Transforms/Utils/LowerSwitch.cpp
+++ llvm/lib/Transforms/Utils/LowerSwitch.cpp
@@ -370,7 +370,9 @@
   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
@@ -431,7 +433,7 @@
 
   if (DefaultIsUnreachableFromSwitch) {
     DenseMap<BasicBlock *, APInt> Popularity;
-    APInt MaxPop(SignedZero);
+    APInt MaxPop(UnsignedZero);
     BasicBlock *PopSucc = nullptr;
 
     APInt SignedMax = APInt::getSignedMaxValue(BitWidth);
@@ -457,11 +459,11 @@
       }
 
       // 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;
       }
@@ -486,8 +488,6 @@
 
     // 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; });
@@ -498,8 +498,9 @@
       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)
-        PopSucc->removePredecessor(OrigBlock);
+      if (!MaxPop.isZero())
+        for (APInt I(UnsignedZero); I.ult(MaxPop - 1); ++I)
+          PopSucc->removePredecessor(OrigBlock);
       return;
     }
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D140747.488721.patch
Type: text/x-patch
Size: 3492 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230112/4af4e211/attachment.bin>