[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 ↗

(On Diff #152039)

Remove the FIXME rather than marking it done.

1405 ↗

(On Diff #152039)

Remvoe the FIXME

1564 ↗

(On Diff #152039)

Space after 'if' before paren

1566 ↗

(On Diff #152039)

Line this up with the InitX on the line above.

1570 ↗

(On Diff #152039)

Space before 'else'

test/Transforms/LoopIdiom/X86/ctlz.ll

118 ↗

(On Diff #152039)

Where are the checks for this test case?

121 ↗

(On Diff #152039)

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

213 ↗

(On Diff #152039)

Remvoe the absolute value

299 ↗

(On Diff #152039)

Remove the absolute value

388 ↗

(On Diff #152039)

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

1569 ↗

(On Diff #152590)

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.

llvm/trunk/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)
+  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))
+  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 =
bool LoopIdiomRecognize::recognizeAndInsertCTLZ() {
   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;
 
-  transformLoopToCountable(PH, CntInst, CntPhi, InitX, DefX->getDebugLoc(),
-                           ZeroCheck, IsCntPhiUsedOutsideLoop);
+  transformLoopToCountable(PH, CntInst, CntPhi, InitX, DefX,
+                           DefX->getDebugLoc(), 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.
 bool LoopIdiomRecognize::recognizePopcount() {

 /// 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)
-    InitXNext = Builder.CreateAShr(InitX,
-                                   ConstantInt::get(InitX->getType(), 1));
+  if (IsCntPhiUsedOutsideLoop) {
+    if (DefX->getOpcode() == Instruction::AShr)
+      InitXNext =
+          Builder.CreateAShr(InitX, ConstantInt::get(InitX->getType(), 1));
+    else if (DefX->getOpcode() == Instruction::LShr)
+      InitXNext =
+          Builder.CreateLShr(InitX, ConstantInt::get(InitX->getType(), 1));
     else
+      llvm_unreachable("Unexpected opcode!");      
+  } else
     InitXNext = InitX;
   CTLZ = createCTLZIntrinsic(Builder, InitXNext, DL, ZeroCheck);
   Count = Builder.CreateSub(
       ConstantInt::get(CTLZ->getType(),
                        CTLZ->getType()->getIntegerBitWidth()),
       CTLZ);
   if (IsCntPhiUsedOutsideLoop) {
     CountPrev = Count;

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

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_zero_check_lshr(int n)
+; {
+;   int i = 0;
+;   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:  %inc.lcssa = phi i32 [ %1, %while.body ]
+; ALL:  %i.0.lcssa = phi i32 [ 0, %entry ], [ %inc.lcssa, %while.end.loopexit ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_zero_check_lshr(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 ], [ 0, %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 [ 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(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_lshr(int n)
+; {
+;   int i = 0;
+;   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:  %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:
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %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)
+; {
+;   int i = i0;
+;   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 = 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:
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %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 i = 0;
+;   while(in >>= 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:  %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
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %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: