Use APInt for threshold in matchShiftULTCondition

Use APInt for the threshold in matchShiftULTCondition, so that the
function works for integers of arbitrary precision (rather than just
those within the range of uint64_t).

Author: hazzlim

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

@@ -1530,7 +1530,7 @@ static Value *matchCondition(BranchInst *BI, BasicBlock *LoopEntry,
 /// the control yields to the loop entry. If the branch matches the behaviour,
 /// the variable involved in the comparison is returned.
 static Value *matchShiftULTCondition(BranchInst *BI, BasicBlock *LoopEntry,
-                                     uint64_t &Threshold) {
+                                     APInt &Threshold) {
   if (!BI || !BI->isConditional())
     return nullptr;
 
@@ -1546,7 +1546,7 @@ static Value *matchShiftULTCondition(BranchInst *BI, BasicBlock *LoopEntry,
   ICmpInst::Predicate Pred = Cond->getPredicate();
 
   if (Pred == ICmpInst::ICMP_ULT && FalseSucc == LoopEntry) {
-    Threshold = CmpConst->getZExtValue();
+    Threshold = CmpConst->getValue();
     return Cond->getOperand(0);
   }
 
@@ -1595,7 +1595,7 @@ static bool detectShiftUntilLessThanIdiom(Loop *CurLoop, const DataLayout &DL,
                                           Intrinsic::ID &IntrinID,
                                           Value *&InitX, Instruction *&CntInst,
                                           PHINode *&CntPhi, Instruction *&DefX,
-                                          uint64_t &Threshold) {
+                                          APInt &Threshold) {
   BasicBlock *LoopEntry;
 
   DefX = nullptr;
@@ -2012,7 +2012,7 @@ bool LoopIdiomRecognize::recognizeShiftUntilLessThan() {
   PHINode *CntPhi = nullptr;
   Instruction *CntInst = nullptr;
 
-  uint64_t LoopThreshold;
+  APInt LoopThreshold;
   if (!detectShiftUntilLessThanIdiom(CurLoop, *DL, IntrinID, InitX, CntInst,
                                      CntPhi, DefX, LoopThreshold))
     return false;
@@ -2041,7 +2041,7 @@ bool LoopIdiomRecognize::recognizeShiftUntilLessThan() {
   if (!PreCondBI)
     return false;
 
-  uint64_t PreLoopThreshold;
+  APInt PreLoopThreshold;
   if (matchShiftULTCondition(PreCondBI, PH, PreLoopThreshold) != InitX ||
       PreLoopThreshold != 2)
     return false;

llvm/test/Transforms/LoopIdiom/AArch64/ctlz.ll

@@ -773,6 +773,37 @@ do.end:                                           ; preds = %do.body
   ret i32 %inc.0
 }
 
+; Check that we correctly bail on analysis when the ult comparison is with a
+; constant that exceeds the (unsigned) range of a 64-bit integer, as we currently
+; only handle loopback condition ult 2 or 4.
+
+define i128 @large_constant(i128 noundef %n) {
+; CHECK-LABEL: define i128 @large_constant(
+; CHECK-SAME: i128 noundef [[N:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[DO_BODY:%.*]]
+; CHECK:       do.body:
+; CHECK-NEXT:    [[N_ADDR_0:%.*]] = phi i128 [ [[N]], [[ENTRY:%.*]] ], [ [[SHR:%.*]], [[DO_BODY]] ]
+; CHECK-NEXT:    [[SHR]] = lshr i128 [[N_ADDR_0]], 1
+; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp ult i128 [[N_ADDR_0]], 18446744073709551616
+; CHECK-NEXT:    br i1 [[CMP_NOT]], label [[DO_END:%.*]], label [[DO_BODY]]
+; CHECK:       do.end:
+; CHECK-NEXT:    [[SHR_LCSSA:%.*]] = phi i128 [ [[SHR]], [[DO_BODY]] ]
+; CHECK-NEXT:    ret i128 [[SHR_LCSSA]]
+;
+entry:
+  br label %do.body
+
+do.body:                                          ; preds = %do.body, %entry
+  %n.addr.0 = phi i128 [ %n, %entry ], [ %shr, %do.body ]
+  %shr = lshr i128 %n.addr.0, 1
+  %cmp.not = icmp ult i128 %n.addr.0, 18446744073709551616
+  br i1 %cmp.not, label %do.end, label %do.body
+
+do.end:                                           ; preds = %do.body
+  ret i128 %shr
+}
+
 
 declare i32 @llvm.abs.i32(i32, i1)
 declare i16 @llvm.abs.i16(i16, i1)