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

[David Green via llvm-commits]

================
@@ -748,93 +708,162 @@ 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)
+  // * Not(1bit)
+  struct SelectLikeInfo {
+    Value *Cond;
+    bool IsAuxiliary;
+    bool IsInverted;
+    unsigned ConditionIdx;
+  };
+
+  std::map<Value *, SelectLikeInfo> SelectInfo;
+
+  auto ProcessSelectInfo = [&SelectInfo](Instruction *I) -> void {
+    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};
+      return;
+    }
+
+    if (match(I, m_Not(m_Value(Cond)))) {
+      SelectInfo[I] = {Cond, true, true, 0};
+      return;
+    }
+
+    // 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};
+      return;
+    }
+
+    // An Or(zext(i1 X), Y) can also be treated like a select, with condition X
+    // and values Y|1 and Y.
+    if (auto *BO = dyn_cast<BinaryOperator>(I)) {
+      if (BO->getType()->isIntegerTy(1) || BO->getOpcode() != Instruction::Or)
+        return;
+
+      for (unsigned Idx = 0; Idx < 2; Idx++) {
+        auto *Op = BO->getOperand(Idx);
+        auto It = SelectInfo.find(Op);
+        if (It != SelectInfo.end() && It->second.IsAuxiliary) {
+          SelectInfo[I] = {It->second.Cond, false, It->second.IsInverted, Idx};
+          break;
+        }
+      }
+    }
+  };
+
+  bool AlreadyProcessed = false;
   BasicBlock::iterator BBIt = BB.begin();
   while (BBIt != BB.end()) {
     Instruction *I = &*BBIt++;
-    if (SelectLike SI = SelectLike::match(I)) {
-      if (!TTI->shouldTreatInstructionLikeSelect(I))
-        continue;
+    if (!AlreadyProcessed)
+      ProcessSelectInfo(I);
+    else
+      AlreadyProcessed = false;
 
-      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;
-        }
+    auto It = SelectInfo.find(I);
+    if (It == SelectInfo.end() || It->second.IsAuxiliary)
+      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;
-        }
+    if (!TTI->shouldTreatInstructionLikeSelect(I))
+      continue;
 
-        // We only allow selects in the same group, not other select-like
-        // instructions.
-        if (!isa<SelectInst>(NI))
-          break;
+    Value *Cond = It->second.Cond;
+    SelectGroup SIGroup{Cond};
+    SIGroup.Selects.emplace_back(I, It->second.IsInverted,
+                                 It->second.ConditionIdx);
 
-        SelectLike NSI = SelectLike::match(NI);
-        if (NSI && SI.getCondition() == NSI.getCondition()) {
-          SIGroup.push_back(NSI);
-        } else if (NSI && match(NSI.getCondition(),
-                                m_Not(m_Specific(SI.getCondition())))) {
-          NSI.setInverted();
-          SIGroup.push_back(NSI);
-        } else
-          break;
+    if (!Cond->getType()->isIntegerTy(1))
----------------
davemgreen wrote:

Can you add a comment to say we do not try to support vector conditions.

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