[LIR] Strengthen the check for recurrence variable in popcnt/CTLZ

[davide]

Summary

Fixes PR33114 (see https://bugs.llvm.org/show_bug.cgi?id=33114 for discussion).

[evstupac]

Thanks for catching this.
The patch looks good, just some inline comments to optimize the code.

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1042

I wonder if you need PhiX as a parameter. It looks like it is unused after the call.
Also it looks like the function does not modify anything, so you don't need references.

test/Transforms/LoopIdiom/pr33114.ll

2

I would mention that the test checks just no assert in comments.
That way we don't need all these IR output checks.

[davide]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1042

Isn't it used in recognizeAndInsertCTLZ ?

1042

Yes, I removed the refs.

[evstupac]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1042

Isn't it used in recognizeAndInsertCTLZ ?

Yes. Right.
I was confused by function name "check". Maybe we could calculate PhiX before the check (outside the function)?
Or do like:

// checks that VarX and DefX create a loop recurrence and return corresponding phi node or NULL otherwise.
static PHINode *checkRecurrence(Value *VarX, Instruction *DefX, BasicBlock *LoopEntry) {
  PhiX = dyn_cast<PHINode>(VarX);
  if (PhiX && PhiX->getParent() == LoopEntry &&
      (PhiX->getOperand(0) == DefX || PhiX->getOperand(1) == DefX))
    return PhiX;
  return nullptr;
}

1142

Can we reuse the new function here for Inst and Inst->getOperand(0)?

1258

And here.

[davide]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1142

This is NFC, I assume, but we're technically changing the check. If you don't mind, I would do that in a follow-up.

[davide]

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1042

I agree it should have a different name. I propose getRecurrenceVar

test/Transforms/LoopIdiom/pr33114.ll

3

Updated, I'll change this before committing.

[davide]

Comments updated.

[evstupac]

LGTM.

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1142

ok.

[davide]

Also, r303704 for the cleanup.

lib/Transforms/Scalar/LoopIdiomRecognize.cpp

   ICmpInst::Predicate Pred = Cond->getPredicate();
   if ((Pred == ICmpInst::ICMP_NE && BI->getSuccessor(0) == LoopEntry) ||
       (Pred == ICmpInst::ICMP_EQ && BI->getSuccessor(1) == LoopEntry))
     return Cond->getOperand(0);
 
   return nullptr;
 }
 
+// Check if the recurrence variable `VarX` is in the right form to form
+// the idiom. Returns the value coerced to a PHINode if so.
+static PHINode *getRecurrenceVar(Value *VarX, Instruction *DefX,
+                                 BasicBlock *LoopEntry) {
+  auto *PhiX = dyn_cast<PHINode>(VarX);
+  if (PhiX && PhiX->getParent() == LoopEntry &&
+      (PhiX->getOperand(0) == DefX || PhiX->getOperand(1) == DefX))
+    return PhiX;
+  return nullptr;
+}
+
 /// Return true iff the idiom is detected in the loop.
 ///
 /// Additionally:
 /// 1) \p CntInst is set to the instruction counting the population bit.
 /// 2) \p CntPhi is set to the corresponding phi node.
 /// 3) \p Var is set to the value whose population bits are being counted.
 ///
 /// The core idiom we are trying to detect is:

         !((SubInst->getOpcode() == Instruction::Sub && Dec->isOne()) ||
           (SubInst->getOpcode() == Instruction::Add &&
            Dec->isAllOnesValue()))) {
       return false;
     }
   }
 
   // step 3: Check the recurrence of variable X
-  {
-    PhiX = dyn_cast<PHINode>(VarX1);
-    if (!PhiX ||
-        (PhiX->getOperand(0) != DefX2 && PhiX->getOperand(1) != DefX2)) {
-      return false;
-    }
-  }
+  PhiX = getRecurrenceVar(VarX1, DefX2, LoopEntry);
+  if (!PhiX)
+    return false;
 
   // step 4: Find the instruction which count the population: cnt2 = cnt1 + 1
   {
     CountInst = nullptr;
     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())
         continue;
 
       PHINode *Phi = dyn_cast<PHINode>(Inst->getOperand(0));
       if (!Phi || Phi->getParent() != LoopEntry)
         continue;
 
       // Check if the result of the instruction is live of the loop.
       bool LiveOutLoop = false;
       for (User *U : Inst->users()) {
         if ((cast<Instruction>(U))->getParent() != LoopEntry) {
           LiveOutLoop = true;

   if (!DefX || DefX->getOpcode() != Instruction::AShr)
     return false;
   if (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 = dyn_cast<PHINode>(VarX);
-  if (!PhiX || (PhiX->getOperand(0) != DefX && PhiX->getOperand(1) != DefX))
+  PhiX = getRecurrenceVar(VarX, DefX, LoopEntry);
+  if (!PhiX)
     return false;
 
   // step 4: Find the instruction which count the CTLZ: 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())
       continue;
 
     PHINode *Phi = dyn_cast<PHINode>(Inst->getOperand(0));
     if (!Phi || Phi->getParent() != LoopEntry)
       continue;
 
     CntInst = Inst;
     CntPhi = Phi;
     break;
   }
   if (!CntInst)

test/Transforms/LoopIdiom/pr33114.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; Check that we're not crashing while looking at the recurrence variable.
+; RUN: opt -S -loop-idiom %s | FileCheck %s
+
+define void @tinkywinky() {
+; CHECK-LABEL: @tinkywinky(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 true, label [[EXIT:%.*]], label [[PH:%.*]]
+; CHECK:       ph:
+; CHECK-NEXT:    [[MYPHI:%.*]] = phi i32 [ 1, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    br label [[IF_END:%.*]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[PATATINO:%.*]] = ashr i32 [[MYPHI]], undef
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i32 [[PATATINO]], 0
+; CHECK-NEXT:    br i1 [[TOBOOL]], label [[EXIT_LOOPEXIT:%.*]], label [[IF_END]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 true, label %exit, label %ph
+
+ph:
+  %myphi = phi i32 [ 1, %entry ]
+  br label %if.end
+
+if.end:
+  %patatino = ashr i32 %myphi, undef
+  %tobool = icmp eq i32 %patatino, 0
+  br i1 %tobool, label %exit, label %if.end
+
+exit:
+  ret void
+}