[LoopIdiomRecognize] Support for loops that use LSHR instruction added.

[ovmold]

Summary

In the 'detectCTLZIdiom' function support for loops that use LSHR instruction instead of ASHR has been added.

The problem is that for the following piece of code no '@llvm.ctlz' instruction has been generated in the resulting test.ll file when compiling as "clang -S -O3 -march=core-avx2 -emit-llvm test.c". The reason for this is that the LSHR instruction is used instead of ASHR in the LLVM IR when we get into the 'detectCTLZIdiom' function.

int lzcnt(int x) {
     int count = 0;
     while (x > 0)  {
          count++;
          x = x >> 1;
     }
    return count;
}

int main() {
    int x;
    scanf("%d", &x);
    int y = lzcnt(x);
    printf("count  = %d\n", y);
    return 0;
}

[craig.topper]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1320

Remove the FIXME rather than marking it done.

1405

Remvoe the FIXME

1564

Space after 'if' before paren

1566

Line this up with the InitX on the line above.

1570

Space before 'else'

test/Transforms/LoopIdiom/X86/ctlz.ll

118

Where are the checks for this test case?

121

Remove the first 3 instructions which calculate an absolute value. And use %n in place of %abs_n in the remaining code.

213

Remvoe the absolute value

299

Remove the absolute value

388

Remove the absolute value.

[ovmold]

Changes in code formatting and test cases according to the reviewer's comments.

[craig.topper]

This def seems to be relative to your previous patch and not relative to the current code in trunk.

[craig.topper]

Oops that should have said "diff" not "def"

[ovmold]

One more attempt: two commits are combined in one patch file.

[bryant]

• Would you mind re-upping this with full context?

• Could you clang-format your changes (not necessarily the entire file)?

• Also, the below condition would not be necessary for LShr:

// 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.
// TODO: Support loops that use lshr and wouldn't need this check.
if (!isKnownNonNegative(InitX, *DL))
  return false;

[ovmold]

• Would you mind re-upping this with full context?

Diff with full context re-upped.

• Could you clang-format your changes (not necessarily the entire file)?

Done.

• Also, the below condition would not be necessary for LShr:
  
  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. // TODO: Support loops that use lshr and wouldn't need this check. if (!isKnownNonNegative(InitX, *DL)) return false;

This has been already done by the first patch. In the current patch it is also present.

[craig.topper]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1573

Can you add

else
  llvm_unreachable("Unexpected opcode!);

It should only be of the two shift opcodes and if for some reason that fails we should error. It also make the code behavior more obvious.

[ovmold]

Can you add

else
  llvm_unreachable("Unexpected opcode!);

It should only be of the two shift opcodes and if for some reason that fails we should error. It also make the code behavior more obvious.

'llvm_unreachable("Unexpected opcode!")' added, code formatting done

[craig.topper]

LGTM

[ovmold]

Who can commit my changes? I have no rights to commit myself.

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

private:
 
   bool runOnNoncountableLoop();
 
   bool recognizePopcount();
   void transformLoopToPopcount(BasicBlock *PreCondBB, Instruction *CntInst,
                                PHINode *CntPhi, Value *Var);
   bool recognizeAndInsertCTLZ();
   void transformLoopToCountable(BasicBlock *PreCondBB, Instruction *CntInst,
-                                PHINode *CntPhi, Value *Var, const DebugLoc DL,
-                                bool ZeroCheck, bool IsCntPhiUsedOutsideLoop);
+                                PHINode *CntPhi, Value *Var, Instruction *DefX, 
+                                const DebugLoc DL, bool ZeroCheck, 
+                                bool IsCntPhiUsedOutsideLoop);
 
   /// @}
 };
 
 class LoopIdiomRecognizeLegacyPass : public LoopPass {
 public:
   static char ID;
 
static bool detectCTLZIdiom(Loop *CurLoop, PHINode *&PhiX,
   // step 1: Check if the loop-back branch is in desirable form.
   if (Value *T = matchCondition(
           dyn_cast<BranchInst>(LoopEntry->getTerminator()), LoopEntry))
     DefX = dyn_cast<Instruction>(T);
   else
     return false;
 
   // step 2: detect instructions corresponding to "x.next = x >> 1"
-  // TODO: Support loops that use LShr.
-  if (!DefX || DefX->getOpcode() != Instruction::AShr)
+  // DONE: Support loops that use LShr.
+  if (!DefX || (DefX->getOpcode() != Instruction::AShr && DefX->getOpcode() != Instruction::LShr))
     return false;
   ConstantInt *Shft = dyn_cast<ConstantInt>(DefX->getOperand(1));
   if (!Shft || !Shft->isOne())
     return false;
   VarX = DefX->getOperand(0);
 
   // step 3: Check the recurrence of variable X
   PhiX = getRecurrenceVar(VarX, DefX, LoopEntry);
bool LoopIdiomRecognize::recognizeAndInsertCTLZ() {
   bool ZeroCheck = false;
   // It is safe to assume Preheader exist as it was checked in
   // parent function RunOnLoop.
   BasicBlock *PH = CurLoop->getLoopPreheader();
   Value *InitX = PhiX->getIncomingValueForBlock(PH);
 
   // 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.
-  // TODO: Support loops that use lshr and wouldn't need this check.
-  if (!isKnownNonNegative(InitX, *DL))
+  // DONE: Support loops that use lshr and wouldn't need this check.
+  if (DefX->getOpcode() == Instruction::AShr && !isKnownNonNegative(InitX, *DL))
     return false;
 
   // If we check X != 0 before entering the loop we don't need a zero
   // check in CTLZ intrinsic, but only if Cnt Phi is not used outside of the
   // loop (if it is used we count CTLZ(X >> 1)).
   if (!IsCntPhiUsedOutsideLoop)
     if (BasicBlock *PreCondBB = PH->getSinglePredecessor())
       if (BranchInst *PreCondBr =
const Value *Args[] =
       {InitX, ZeroCheck ? ConstantInt::getTrue(InitX->getContext())
                         : ConstantInt::getFalse(InitX->getContext())};
   if (CurLoop->getHeader()->size() != 6 &&
       TTI->getIntrinsicCost(Intrinsic::ctlz, InitX->getType(), Args) >
           TargetTransformInfo::TCC_Basic)
     return false;
 
   const DebugLoc DL = DefX->getDebugLoc();
-  transformLoopToCountable(PH, CntInst, CntPhi, InitX, DL, ZeroCheck,
+  transformLoopToCountable(PH, CntInst, CntPhi, InitX, DefX, DL, ZeroCheck,
                            IsCntPhiUsedOutsideLoop);
   return true;
 }
 
 /// Recognizes a population count idiom in a non-countable loop.
 ///
 /// If detected, transforms the relevant code to issue the popcount intrinsic
 /// function call, and returns true; otherwise, returns false.

 /// Use(CountPrev + Cnt0) // Use(CntPhi)
 /// or
 /// Use(Count + Cnt0) // Use(CntInst)
 ///
 /// If LOOP_BODY is empty the loop will be deleted.
 /// If CntInst and DefX are not used in LOOP_BODY they will be removed.
 void LoopIdiomRecognize::transformLoopToCountable(
     BasicBlock *Preheader, Instruction *CntInst, PHINode *CntPhi, Value *InitX,
-    const DebugLoc DL, bool ZeroCheck, bool IsCntPhiUsedOutsideLoop) {
+    Instruction *DefX, const DebugLoc DL, bool ZeroCheck, bool IsCntPhiUsedOutsideLoop) {
   BranchInst *PreheaderBr = cast<BranchInst>(Preheader->getTerminator());
 
   // Step 1: Insert the CTLZ instruction at the end of the preheader block
   //   Count = BitWidth - CTLZ(InitX);
   // If there are uses of CntPhi create:
   //   CountPrev = BitWidth - CTLZ(InitX >> 1);
   IRBuilder<> Builder(PreheaderBr);
   Builder.SetCurrentDebugLocation(DL);
   Value *CTLZ, *Count, *CountPrev, *NewCount, *InitXNext;
 
-  if (IsCntPhiUsedOutsideLoop)
+  if (IsCntPhiUsedOutsideLoop){
+    if(DefX->getOpcode() == Instruction::AShr)
       InitXNext = Builder.CreateAShr(InitX,
                                    ConstantInt::get(InitX->getType(), 1));
-  else
+    else if(DefX->getOpcode() == Instruction::LShr)
+      InitXNext = Builder.CreateLShr(InitX,
+                                   ConstantInt::get(InitX->getType(), 1));
+  }else
     InitXNext = InitX;
   CTLZ = createCTLZIntrinsic(Builder, InitXNext, DL, ZeroCheck);
   Count = Builder.CreateSub(
       ConstantInt::get(CTLZ->getType(),
                        CTLZ->getType()->getIntegerBitWidth()),
       CTLZ);
   if (IsCntPhiUsedOutsideLoop) {
     CountPrev = Count;
     Count = Builder.CreateAdd(
         CountPrev,
         ConstantInt::get(CountPrev->getType(), 1));

test/Transforms/LoopIdiom/X86/ctlz.ll

 while.end.loopexit:                               ; preds = %while.body
   br label %while.end
 
 while.end:                                        ; preds = %while.end.loopexit, %entry
   %i.0.lcssa = phi i32 [ 0, %entry ], [ %inc, %while.end.loopexit ]
   ret i32 %i.0.lcssa
 }
 
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_zero_check_lshr(i32 %n) {
+entry:
+  %c = icmp sgt i32 %n, 0
+  %negn = sub nsw i32 0, %n
+  %abs_n = select i1 %c, i32 %n, i32 %negn
+  %tobool4 = icmp eq i32 %abs_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 ], [ 0, %while.body.preheader ]
+  %n.addr.05 = phi i32 [ %shr, %while.body ], [ %abs_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 [ 0, %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(int n)
 ; {
 ;   n = n >= 0 ? n : -n;
 ;   int i = 0;

 while.end:                                        ; preds = %while.cond
   ret i32 %i.0
 }
 
 ; Recognize CTLZ builtin pattern.
 ; Here it will replace the loop -
 ; assume builtin is always profitable.
 ;
+; int ctlz_lshr(int n)
+; {
+;   n = n >= 0 ? n : -n;
+;   int i = 0;
+;   while(n >>= 1) {
+;     i++;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = lshr i32 %abs_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:  %i.0.lcssa = phi i32 [ %2, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_lshr(i32 %n) {
+entry:
+  %c = icmp sgt i32 %n, 0
+  %negn = sub nsw i32 0, %n
+  %abs_n = select i1 %c, i32 %n, i32 %negn
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %abs_n, %entry ], [ %shr, %while.cond ]
+  %i.0 = phi i32 [ 0, %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
+}
+
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
 ; int ctlz_add(int n, int i0)
 ; {
 ;   n = n >= 0 ? n : -n;
 ;   int i = i0;
 ;   while(n >>= 1) {
 ;     i++;
 ;   }
 ;   return i;

 while.end:                                        ; preds = %while.cond
   ret i32 %i.0
 }
 
 ; Recognize CTLZ builtin pattern.
 ; Here it will replace the loop -
 ; assume builtin is always profitable.
 ;
+; int ctlz_add_lshr(int n, int i0)
+; {
+;   n = n >= 0 ? n : -n;
+;   int i = i0;
+;   while(n >>= 1) {
+;     i++;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = lshr i32 %abs_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 = add i32 %2, %i0
+; ALL:  %i.0.lcssa = phi i32 [ %4, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+;
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_add_lshr(i32 %n, i32 %i0) {
+entry:
+  %c = icmp sgt i32 %n, 0
+  %negn = sub nsw i32 0, %n
+  %abs_n = select i1 %c, i32 %n, i32 %negn
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %abs_n, %entry ], [ %shr, %while.cond ]
+  %i.0 = phi i32 [ %i0, %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
+}
+
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
 ; int ctlz_sext(short in)
 ; {
 ;   int n = in;
 ;   if (in < 0)
 ;     n = -n;
 ;   int i = 0;
 ;   while(n >>= 1) {
 ;     i++;
while.cond:                                       ; preds = %while.cond, %entry
   %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
 }
 
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int ctlz_sext_lshr(short in)
+; {
+;   int n = in;
+;   if (in < 0)
+;     n = -n;
+;   int i = 0;
+;   while(n >>= 1) {
+;     i++;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = lshr i32 %abs_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:  %i.0.lcssa = phi i32 [ %2, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_sext_lshr(i16 %in) {
+entry:
+  %n = sext i16 %in to i32
+  %c = icmp sgt i16 %in, 0
+  %negn = sub nsw i32 0, %n
+  %abs_n = select i1 %c, i32 %n, i32 %negn
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %abs_n, %entry ], [ %shr, %while.cond ]
+  %i.0 = phi i32 [ 0, %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
+}
+
 ; This loop contains a volatile store. If x is initially negative,
 ; the code will be an infinite loop because the ashr will eventually produce
 ; all ones and continue doing so. This prevents the loop from terminating. If
 ; we convert this to a countable loop using ctlz that loop will only run 32
 ; times. This is different than the infinite number of times of the original.
 define i32 @foo(i32 %x) {
 ; LZCNT-LABEL: @foo(
 ; LZCNT-NEXT:  entry: