[DFAJumpThreading] Extends the bitwidth of state from uint64_t to APInt (#78134)

Fixes #78059

Author: XChy

llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp

@@ -307,7 +307,7 @@ void unfold(DomTreeUpdater *DTU, SelectInstToUnfold SIToUnfold,
 
 struct ClonedBlock {
   BasicBlock *BB;
-  uint64_t State; ///< \p State corresponds to the next value of a switch stmnt.
+  APInt State; ///< \p State corresponds to the next value of a switch stmnt.
 };
 
 typedef std::deque<BasicBlock *> PathType;
@@ -344,9 +344,9 @@ inline raw_ostream &operator<<(raw_ostream &OS, const PathType &Path) {
 /// exit state, and the block that determines the next state.
 struct ThreadingPath {
   /// Exit value is DFA's exit state for the given path.
-  uint64_t getExitValue() const { return ExitVal; }
+  APInt getExitValue() const { return ExitVal; }
   void setExitValue(const ConstantInt *V) {
-    ExitVal = V->getZExtValue();
+    ExitVal = V->getValue();
     IsExitValSet = true;
   }
   bool isExitValueSet() const { return IsExitValSet; }
@@ -365,7 +365,7 @@ struct ThreadingPath {
 
 private:
   PathType Path;
-  uint64_t ExitVal;
+  APInt ExitVal;
   const BasicBlock *DBB = nullptr;
   bool IsExitValSet = false;
 };
@@ -744,7 +744,7 @@ struct TransformDFA {
 
     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
       PathType PathBBs = TPath.getPath();
-      uint64_t NextState = TPath.getExitValue();
+      APInt NextState = TPath.getExitValue();
       const BasicBlock *Determinator = TPath.getDeterminatorBB();
 
       // Update Metrics for the Switch block, this is always cloned
@@ -901,7 +901,7 @@ struct TransformDFA {
                       DuplicateBlockMap &DuplicateMap,
                       SmallSet<BasicBlock *, 16> &BlocksToClean,
                       DomTreeUpdater *DTU) {
-    uint64_t NextState = Path.getExitValue();
+    APInt NextState = Path.getExitValue();
     const BasicBlock *Determinator = Path.getDeterminatorBB();
     PathType PathBBs = Path.getPath();
 
@@ -993,13 +993,14 @@ struct TransformDFA {
   /// This function also includes updating phi nodes in the successors of the
   /// BB, and remapping uses that were defined locally in the cloned BB.
   BasicBlock *cloneBlockAndUpdatePredecessor(BasicBlock *BB, BasicBlock *PrevBB,
-                                             uint64_t NextState,
+                                             const APInt &NextState,
                                              DuplicateBlockMap &DuplicateMap,
                                              DefMap &NewDefs,
                                              DomTreeUpdater *DTU) {
     ValueToValueMapTy VMap;
     BasicBlock *NewBB = CloneBasicBlock(
-        BB, VMap, ".jt" + std::to_string(NextState), BB->getParent());
+        BB, VMap, ".jt" + std::to_string(NextState.getLimitedValue()),
+        BB->getParent());
     NewBB->moveAfter(BB);
     NumCloned++;
 
@@ -1034,7 +1035,7 @@ struct TransformDFA {
   /// This means creating a new incoming value from NewBB with the new
   /// instruction wherever there is an incoming value from BB.
   void updateSuccessorPhis(BasicBlock *BB, BasicBlock *ClonedBB,
-                           uint64_t NextState, ValueToValueMapTy &VMap,
+                           const APInt &NextState, ValueToValueMapTy &VMap,
                            DuplicateBlockMap &DuplicateMap) {
     std::vector<BasicBlock *> BlocksToUpdate;
 
@@ -1144,7 +1145,7 @@ struct TransformDFA {
   void updateLastSuccessor(ThreadingPath &TPath,
                            DuplicateBlockMap &DuplicateMap,
                            DomTreeUpdater *DTU) {
-    uint64_t NextState = TPath.getExitValue();
+    APInt NextState = TPath.getExitValue();
     BasicBlock *BB = TPath.getPath().back();
     BasicBlock *LastBlock = getClonedBB(BB, NextState, DuplicateMap);
 
@@ -1198,7 +1199,7 @@ struct TransformDFA {
 
   /// Checks if BB was already cloned for a particular next state value. If it
   /// was then it returns this cloned block, and otherwise null.
-  BasicBlock *getClonedBB(BasicBlock *BB, uint64_t NextState,
+  BasicBlock *getClonedBB(BasicBlock *BB, const APInt &NextState,
                           DuplicateBlockMap &DuplicateMap) {
     CloneList ClonedBBs = DuplicateMap[BB];
 
@@ -1212,10 +1213,10 @@ struct TransformDFA {
 
   /// Helper to get the successor corresponding to a particular case value for
   /// a switch statement.
-  BasicBlock *getNextCaseSuccessor(SwitchInst *Switch, uint64_t NextState) {
+  BasicBlock *getNextCaseSuccessor(SwitchInst *Switch, const APInt &NextState) {
     BasicBlock *NextCase = nullptr;
     for (auto Case : Switch->cases()) {
-      if (Case.getCaseValue()->getZExtValue() == NextState) {
+      if (Case.getCaseValue()->getValue() == NextState) {
         NextCase = Case.getCaseSuccessor();
         break;
       }

llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll

@@ -12,8 +12,8 @@ define i32 @test1(i32 %num) {
 ; CHECK-NEXT:    [[COUNT:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_INC:%.*]] ]
 ; CHECK-NEXT:    [[STATE:%.*]] = phi i32 [ 1, [[ENTRY]] ], [ poison, [[FOR_INC]] ]
 ; CHECK-NEXT:    switch i32 [[STATE]], label [[FOR_INC_JT1:%.*]] [
-; CHECK-NEXT:    i32 1, label [[CASE1:%.*]]
-; CHECK-NEXT:    i32 2, label [[CASE2:%.*]]
+; CHECK-NEXT:      i32 1, label [[CASE1:%.*]]
+; CHECK-NEXT:      i32 2, label [[CASE2:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       for.body.jt2:
 ; CHECK-NEXT:    [[COUNT_JT2:%.*]] = phi i32 [ [[INC_JT2:%.*]], [[FOR_INC_JT2:%.*]] ]
@@ -127,11 +127,11 @@ define i32 @test2(i32 %init) {
 ; CHECK:       loop.3:
 ; CHECK-NEXT:    [[STATE:%.*]] = phi i32 [ [[STATE_2]], [[LOOP_2]] ]
 ; CHECK-NEXT:    switch i32 [[STATE]], label [[INFLOOP_I:%.*]] [
-; CHECK-NEXT:    i32 2, label [[CASE2:%.*]]
-; CHECK-NEXT:    i32 3, label [[CASE3:%.*]]
-; CHECK-NEXT:    i32 4, label [[CASE4:%.*]]
-; CHECK-NEXT:    i32 0, label [[CASE0:%.*]]
-; CHECK-NEXT:    i32 1, label [[CASE1:%.*]]
+; CHECK-NEXT:      i32 2, label [[CASE2:%.*]]
+; CHECK-NEXT:      i32 3, label [[CASE3:%.*]]
+; CHECK-NEXT:      i32 4, label [[CASE4:%.*]]
+; CHECK-NEXT:      i32 0, label [[CASE0:%.*]]
+; CHECK-NEXT:      i32 1, label [[CASE1:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       loop.3.jt2:
 ; CHECK-NEXT:    [[STATE_JT2:%.*]] = phi i32 [ [[STATE_2_JT2]], [[LOOP_2_JT2]] ]
@@ -232,3 +232,38 @@ infloop.i:
 exit:
   ret i32 0
 }
+
+define void @pr78059_bitwidth(i1 %c) {
+; CHECK-LABEL: @pr78059_bitwidth(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[DOTSPLIT_PREHEADER:%.*]], label [[DOTSPLIT_PREHEADER]]
+; CHECK:       .split.preheader:
+; CHECK-NEXT:    br label [[DOTSPLIT:%.*]]
+; CHECK:       .split:
+; CHECK-NEXT:    [[TMP0:%.*]] = phi i128 [ 0, [[DOTSPLIT_PREHEADER]] ]
+; CHECK-NEXT:    switch i128 [[TMP0]], label [[END:%.*]] [
+; CHECK-NEXT:      i128 -1, label [[END]]
+; CHECK-NEXT:      i128 0, label [[DOTSPLIT_JT18446744073709551615:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       .split.jt18446744073709551615:
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i128 [ -1, [[DOTSPLIT]] ]
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %c, label %.split.preheader, label %.split.preheader
+
+.split.preheader:
+  br label %.split
+
+.split:
+  %0 = phi i128 [ 0, %.split.preheader ], [ -1, %.split ]
+  switch i128 %0, label %end [
+  i128 -1, label %end
+  i128 0, label %.split
+  ]
+
+end:
+  ret void
+}