[llvm] [SelectOpt] Refactor to prepare for support more select-like operations (PR #115745)

[Igor Kirillov via llvm-commits]

================
@@ -749,93 +717,160 @@ void SelectOptimizeImpl::convertProfitableSIGroups(SelectGroups &ProfSIGroups) {
       FT = FalseBlock;
     }
     IRBuilder<> IB(SI.getI());
-    auto *CondFr = IB.CreateFreeze(SI.getCondition(),
-                                   SI.getCondition()->getName() + ".frozen");
+    auto *CondFr =
+        IB.CreateFreeze(ASI.Condition, ASI.Condition->getName() + ".frozen");
 
-    SmallPtrSet<const Instruction *, 2> INS;
-    for (auto SI : ASI)
-      INS.insert(SI.getI());
+    SmallDenseMap<Instruction *, std::pair<Value *, Value *>, 2> INS;
 
     // Use reverse iterator because later select may use the value of the
     // earlier select, and we need to propagate value through earlier select
     // to get the PHI operand.
-    for (auto It = ASI.rbegin(); It != ASI.rend(); ++It) {
-      SelectLike SI = *It;
+    InsertionPoint = EndBlock->begin();
+    for (SelectLike &SI : ASI.Selects) {
       // The select itself is replaced with a PHI Node.
       PHINode *PN = PHINode::Create(SI.getType(), 2, "");
-      PN->insertBefore(EndBlock->begin());
+      PN->insertBefore(InsertionPoint);
       PN->takeName(SI.getI());
-      PN->addIncoming(getTrueOrFalseValue(SI, true, INS, IB), TrueBlock);
-      PN->addIncoming(getTrueOrFalseValue(SI, false, INS, IB), FalseBlock);
-      PN->setDebugLoc(SI.getI()->getDebugLoc());
+      // Current instruction might be a condition of some other group, so we
+      // need to replace it there to avoid dangling pointer
+      if (PN->getType()->isIntegerTy(1)) {
+        for (auto &SG : ProfSIGroups) {
+          if (SG.Condition == SI.getI())
+            SG.Condition = PN;
+        }
+      }
       SI.getI()->replaceAllUsesWith(PN);
-      INS.erase(SI.getI());
+      auto *TV = getTrueOrFalseValue(SI, true, INS, TrueBlock);
+      auto *FV = getTrueOrFalseValue(SI, false, INS, FalseBlock);
+      INS[PN] = {TV, FV};
+      PN->addIncoming(TV, TrueBlock);
+      PN->addIncoming(FV, FalseBlock);
+      PN->setDebugLoc(SI.getI()->getDebugLoc());
       ++NumSelectsConverted;
     }
     IB.CreateCondBr(CondFr, TT, FT, SI.getI());
 
     // Remove the old select instructions, now that they are not longer used.
-    for (auto SI : ASI)
+    for (SelectLike &SI : ASI.Selects)
       SI.getI()->eraseFromParent();
   }
 }
 
 void SelectOptimizeImpl::collectSelectGroups(BasicBlock &BB,
                                              SelectGroups &SIGroups) {
+  // Represents something that can be considered as select instruction.
+  // Auxiliary instruction are instructions that depends on a condition and have
+  // zero or some constant value on True/False branch, such as:
+  // * ZExt(1bit), SExt(1bit)
+  // * Not(1bit)
+  // * AShr(Xbit), X-1, LShr(XBit), X-1, where there is a condition like Xbit <=
+  // 0 somewhere above in BB
+  struct SelectLikeInfo {
+    Value *Cond;
+    bool IsAuxiliary;
+    bool IsInverted;
+    unsigned ConditionIdx;
+  };
+
+  SmallPtrSet<Instruction *, 2> SeenCmp;
+  std::map<Value *, SelectLikeInfo> SelectInfo;
+
   BasicBlock::iterator BBIt = BB.begin();
   while (BBIt != BB.end()) {
     Instruction *I = &*BBIt++;
-    if (SelectLike SI = SelectLike::match(I)) {
-      if (!TTI->shouldTreatInstructionLikeSelect(I))
+    if (auto *Cmp = dyn_cast<CmpInst>(I)) {
+      SeenCmp.insert(Cmp);
+      continue;
+    }
+
+    Value *Cond;
+    if (match(I, m_OneUse(m_ZExt(m_Value(Cond)))) &&
+        Cond->getType()->isIntegerTy(1)) {
+      bool Inverted = match(Cond, m_Not(m_Value(Cond)));
+      SelectInfo[I] = {Cond, true, Inverted, 0};
+      continue;
+    }
+
+    if (match(I, m_Not(m_Value(Cond)))) {
+      SelectInfo[I] = {Cond, true, true, 0};
+      continue;
+    }
+
+    // Select instruction are what we are usually looking for.
+    if (match(I, m_Select(m_Value(Cond), m_Value(), m_Value()))) {
+      bool Inverted = match(Cond, m_Not(m_Value(Cond)));
+      SelectInfo[I] = {Cond, false, Inverted, 0};
+      continue;
+    }
+
+    // An Or(zext(i1 X), Y) can also be treated like a select, with condition
+    if (auto *BO = dyn_cast<BinaryOperator>(I)) {
+      if (BO->getType()->isIntegerTy(1) || BO->getOpcode() != Instruction::Or)
         continue;
 
-      SelectGroup SIGroup;
-      SIGroup.push_back(SI);
-      while (BBIt != BB.end()) {
-        Instruction *NI = &*BBIt;
-        // Debug/pseudo instructions should be skipped and not prevent the
-        // formation of a select group.
-        if (NI->isDebugOrPseudoInst()) {
-          ++BBIt;
-          continue;
+      for (unsigned Idx = 0; Idx < 2; Idx++) {
+        auto *Op = BO->getOperand(Idx);
+        if (SelectInfo.count(Op) && SelectInfo[Op].IsAuxiliary) {
+          Cond = SelectInfo[Op].Cond;
+          bool Inverted = SelectInfo[Op].IsInverted;
+          SelectInfo[I] = {Cond, false, Inverted, Idx};
+          break;
         }
+      }
+      continue;
+    }
+  }
 
-        // Skip not(select(..)), if the not is part of the same select group
-        if (match(NI, m_Not(m_Specific(SI.getCondition())))) {
-          ++BBIt;
-          continue;
-        }
+  BBIt = BB.begin();
+  while (BBIt != BB.end()) {
----------------
igogo-x86 wrote:

Something like this?

https://github.com/llvm/llvm-project/pull/115745