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

[Igor Kirillov via llvm-commits]

================
@@ -748,93 +707,167 @@ 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;
+  };
+
+  DenseMap<Value *, SelectLikeInfo> SelectInfo;
+
+  // Check if the instruction is SelectLike or might be part of SelectLike
+  // expression, put information into SelectInfo and return the iterator to the
+  // inserted position.
+  auto ProcessSelectInfo = [&SelectInfo](Instruction *I) {
+    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)));
+      return SelectInfo.insert({I, {Cond, true, Inverted, 0}}).first;
+    }
+
+    if (match(I, m_Not(m_Value(Cond)))) {
+      return SelectInfo.insert({I, {Cond, true, true, 0}}).first;
+    }
+
+    // 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)));
+      return SelectInfo.insert({I, {Cond, false, Inverted, 0}}).first;
+    }
+
+    // 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 SelectInfo.end();
+
+      for (unsigned Idx = 0; Idx < 2; Idx++) {
+        auto *Op = BO->getOperand(Idx);
+        auto It = SelectInfo.find(Op);
+        if (It != SelectInfo.end() && It->second.IsAuxiliary) {
+          Cond = It->second.Cond;
+          bool Inverted = It->second.IsInverted;
+          return SelectInfo.insert({I, {Cond, false, Inverted, Idx}}).first;
+        }
+      }
+    }
+    return SelectInfo.end();
+  };
+
+  bool AlreadyProcessed = false;
   BasicBlock::iterator BBIt = BB.begin();
+  DenseMap<Value *, SelectLikeInfo>::iterator It;
   while (BBIt != BB.end()) {
     Instruction *I = &*BBIt++;
-    if (SelectLike SI = SelectLike::match(I)) {
-      if (!TTI->shouldTreatInstructionLikeSelect(I))
-        continue;
+    if (I->isDebugOrPseudoInst())
+      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;
-        }
+    if (!AlreadyProcessed)
+      It = ProcessSelectInfo(I);
+    else
+      AlreadyProcessed = false;
 
-        // 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 (It == SelectInfo.end() || It->second.IsAuxiliary)
+      continue;
+
+    if (!TTI->shouldTreatInstructionLikeSelect(I))
+      continue;
+
+    Value *Cond = It->second.Cond;
+    // Vector conditions are not supported.
+    if (!Cond->getType()->isIntegerTy(1))
+      continue;
+
+    SelectGroup SIGroup = {Cond, {}};
----------------
igogo-x86 wrote:

@fhahn I made it this way. Round brackets won't work unless I create a constructor and it seems like overkill for such a simple structure

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