[LoopIdiomRecognize] Teach detectShiftUntilZeroIdiom to recognize loops where the counter is decrementing.

[craig.topper]

Summary

This adds support for loops like

unsigned clz(unsigned x) {
    unsigned w = sizeof (x) * 8;
    while (x) {
        w--;
        x >>= 1;
    }

    return w;
}

and

unsigned clz(unsigned x) {
    unsigned w = sizeof (x) * 8 - 1;
    while (x >>= 1) {
        w--;
    }

    return w;
}

To support these we look for add x, -1 as well as add x, 1 that
we already matched. If the value was -1 we need to subtract from
the initial counter value instead of adding to it.

Fixes PR48404.

[lebedev.ri]

Seems to look ok in principle, but needs some negative tests too.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1484–1489

Preexisting problem, but this is kinda broken.
If we have a second add in the loop,
and it happens to be before the one we are looking for,
we'll fail to detect the idiom.
The detection should not involve linear scan over instructions.

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

531

Precommit this + use update script?

[craig.topper]

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

531

Should use the script on the whole file?

[craig.topper]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1484–1489

All the fails are continues. So wouldn’t the second add also have to be a recurrence with 1/-1?

[lebedev.ri]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1484–1489

I guess, but still, this is broken:
https://godbolt.org/z/fzoo61

[lebedev.ri]

I.e. my comment is - given that nowadays -indvars will run right afterwards,
should we really care about trying to do it's job.
Should we not *just* transform the loop into countable,
and let it deal with cleaning the old IV's?

[craig.topper]

Rebase after pre-committing tests

[lebedev.ri]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1756–1767

I think this might be marginally less awkward to write this as:

[lebedev.ri]

LGTM, thank you.
I've checked, and alive2 is also happy with this.
As discussed, once D91038 patchset makes progress, i'd like to either subsume this functionality, or at least rip out indvar canonicalization.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

 //===- LoopIdiomRecognize.cpp - Loop idiom recognition --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
static bool detectShiftUntilZeroIdiom(Loop *CurLoop, const DataLayout &DL,
   InitX = PhiX->getIncomingValueForBlock(CurLoop->getLoopPreheader());
 
   // Make sure the initial value can't be negative otherwise the ashr in the
   // loop might never reach zero which would make the loop infinite.
   if (DefX->getOpcode() == Instruction::AShr && !isKnownNonNegative(InitX, DL))
     return false;
 
   // step 4: Find the instruction which count the CTLZ: cnt.next = cnt + 1
+  //         or cnt.next = cnt + -1.
   // TODO: We can skip the step. If loop trip count is known (CTLZ),
   //       then all uses of "cnt.next" could be optimized to the trip count
   //       plus "cnt0". Currently it is not optimized.
   //       This step could be used to detect POPCNT instruction:
   //       cnt.next = cnt + (x.next & 1)
   for (BasicBlock::iterator Iter = LoopEntry->getFirstNonPHI()->getIterator(),
                             IterE = LoopEntry->end();
        Iter != IterE; Iter++) {
     Instruction *Inst = &*Iter;
     if (Inst->getOpcode() != Instruction::Add)
       continue;
 
     ConstantInt *Inc = dyn_cast<ConstantInt>(Inst->getOperand(1));
-    if (!Inc || !Inc->isOne())
+    if (!Inc || (!Inc->isOne() && !Inc->isMinusOne()))
       continue;
 
     PHINode *Phi = getRecurrenceVar(Inst->getOperand(0), Inst, LoopEntry);
     if (!Phi)
       continue;
 
     CntInst = Inst;
     CntPhi = Phi;
void LoopIdiomRecognize::transformLoopToCountable(
   if (IsCntPhiUsedOutsideLoop) {
     NewCount = Count;
     Count = Builder.CreateAdd(Count, ConstantInt::get(Count->getType(), 1));
   }
 
   NewCount = Builder.CreateZExtOrTrunc(NewCount,
                                        cast<IntegerType>(CntInst->getType()));
 
-  // If the counter's initial value is not zero, insert Add Inst.
   Value *CntInitVal = CntPhi->getIncomingValueForBlock(Preheader);
+  if (cast<ConstantInt>(CntInst->getOperand(1))->isOne()) {
+    // If the counter was being incremented in the loop, add NewCount to the
+    // counter's initial value, but only if the initial value is not zero.
     ConstantInt *InitConst = dyn_cast<ConstantInt>(CntInitVal);
     if (!InitConst || !InitConst->isZero())
       NewCount = Builder.CreateAdd(NewCount, CntInitVal);
+  } else {
+    // If the count was being decremented in the loop, subract NewCount from the
+    // counter's initial value.
+    NewCount = Builder.CreateSub(CntInitVal, NewCount);
+  }
 
   // Step 2: Insert new IV and loop condition:
   // loop:
   //   ...
   //   PhiCount = PHI [Count, Dec]
   //   ...
   //   Dec = PhiCount - 1
   //   ...
   //   Br: loop if (Dec != 0)

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

while.cond:                                       ; preds = %while.cond, %entry
   %shr = lshr i32 %n.addr.0, 1
   %tobool = icmp eq i32 %shr, 0
   %inc = add nsw i32 %i.0, 1
   br i1 %tobool, label %while.end, label %while.cond
 
 while.end:                                        ; preds = %while.cond
   ret i32 %inc
 }
+
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int ctlz_decrement(int n)
+; {
+;   int i = 32;
+;   while(n) {
+;     n >>= 1;
+;     i--;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = call i32 @llvm.ctlz.i32(i32 %n, i1 true)
+; ALL-NEXT:  %1 = sub i32 32, %0
+; ALL-NEXT:  %2 = sub i32 32, %1
+; ALL:  %inc.lcssa = phi i32 [ %2, %while.body ]
+; ALL:  %i.0.lcssa = phi i32 [ 32, %entry ], [ %inc.lcssa, %while.end.loopexit ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_decrement(i32 %n) {
+entry:
+  %tobool4 = icmp eq i32 %n, 0
+  br i1 %tobool4, label %while.end, label %while.body.preheader
+
+while.body.preheader:                             ; preds = %entry
+  br label %while.body
+
+while.body:                                       ; preds = %while.body.preheader, %while.body
+  %i.06 = phi i32 [ %inc, %while.body ], [ 32, %while.body.preheader ]
+  %n.addr.05 = phi i32 [ %shr, %while.body ], [ %n, %while.body.preheader ]
+  %shr = lshr i32 %n.addr.05, 1
+  %inc = add nsw i32 %i.06, -1
+  %tobool = icmp eq i32 %shr, 0
+  br i1 %tobool, label %while.end.loopexit, label %while.body
+
+while.end.loopexit:                               ; preds = %while.body
+  br label %while.end
+
+while.end:                                        ; preds = %while.end.loopexit, %entry
+  %i.0.lcssa = phi i32 [ 32, %entry ], [ %inc, %while.end.loopexit ]
+  ret i32 %i.0.lcssa
+}
+
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int ctlz_lshr_decrement(int n)
+; {
+;   int i = 31;
+;   while(n >>= 1) {
+;     i--;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = lshr i32 %n, 1
+; ALL-NEXT:  %1 = call i32 @llvm.ctlz.i32(i32 %0, i1 false)
+; ALL-NEXT:  %2 = sub i32 32, %1
+; ALL-NEXT:  %3 = add i32 %2, 1
+; ALL-NEXT:  %4 = sub i32 31, %2
+; ALL:  %i.0.lcssa = phi i32 [ %4, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_lshr_decrement(i32 %n) {
+entry:
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %n, %entry ], [ %shr, %while.cond ]
+  %i.0 = phi i32 [ 31, %entry ], [ %inc, %while.cond ]
+  %shr = lshr i32 %n.addr.0, 1
+  %tobool = icmp eq i32 %shr, 0
+  %inc = add nsw i32 %i.0, -1
+  br i1 %tobool, label %while.end, label %while.cond
+
+while.end:                                        ; preds = %while.cond
+  ret i32 %i.0
+}

llvm/test/Transforms/LoopIdiom/X86/cttz.ll

while.cond:                                       ; preds = %while.cond, %entry
   %tobool = icmp eq i32 %shl, 0
   %inc = add nsw i32 %i.0, 1
   br i1 %tobool, label %while.end, label %while.cond
 
 while.end:                                        ; preds = %while.cond
   ret i32 %i.0
 }
 
+; Recognize CTTZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int ctlz_decrement(int n)
+; {
+;   int i = 32;
+;   while(n) {
+;     n <<= 1;
+;     i--;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = call i32 @llvm.cttz.i32(i32 %n, i1 true)
+; ALL-NEXT:  %1 = sub i32 32, %0
+; ALL-NEXT:  %2 = sub i32 32, %1
+; ALL:  %inc.lcssa = phi i32 [ %2, %while.body ]
+; ALL:  %i.0.lcssa = phi i32 [ 32, %entry ], [ %inc.lcssa, %while.end.loopexit ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @cttz_decrement(i32 %n) {
+entry:
+  %tobool4 = icmp eq i32 %n, 0
+  br i1 %tobool4, label %while.end, label %while.body.preheader
+
+while.body.preheader:                             ; preds = %entry
+  br label %while.body
+
+while.body:                                       ; preds = %while.body.preheader, %while.body
+  %i.06 = phi i32 [ %inc, %while.body ], [ 32, %while.body.preheader ]
+  %n.addr.05 = phi i32 [ %shl, %while.body ], [ %n, %while.body.preheader ]
+  %shl = shl i32 %n.addr.05, 1
+  %inc = add nsw i32 %i.06, -1
+  %tobool = icmp eq i32 %shl, 0
+  br i1 %tobool, label %while.end.loopexit, label %while.body
+
+while.end.loopexit:                               ; preds = %while.body
+  br label %while.end
+
+while.end:                                        ; preds = %while.end.loopexit, %entry
+  %i.0.lcssa = phi i32 [ 32, %entry ], [ %inc, %while.end.loopexit ]
+  ret i32 %i.0.lcssa
+}
+
+; Recognize CTTZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int cttz_shl_decrement(int n)
+; {
+;   int i = 31;
+;   while(n <<= 1) {
+;     i--;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = shl i32 %n, 1
+; ALL-NEXT:  %1 = call i32 @llvm.cttz.i32(i32 %0, i1 false)
+; ALL-NEXT:  %2 = sub i32 32, %1
+; ALL-NEXT:  %3 = add i32 %2, 1
+; ALL-NEXT:  %4 = sub i32 31, %2
+; ALL:  %i.0.lcssa = phi i32 [ %4, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @cttz_shl_decrement(i32 %n) {
+entry:
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %n, %entry ], [ %shl, %while.cond ]
+  %i.0 = phi i32 [ 31, %entry ], [ %inc, %while.cond ]
+  %shl = shl i32 %n.addr.0, 1
+  %tobool = icmp eq i32 %shl, 0
+  %inc = add nsw i32 %i.0, -1
+  br i1 %tobool, label %while.end, label %while.cond
+
+while.end:                                        ; preds = %while.cond
+  ret i32 %i.0
+}