[SelectOpt] Support ADD and SUB with zext operands. (#115489)

Extend the support for implicit selects in the form of OR with a ZExt
operand to support ADD and SUB binops as well. They similarly can form
implicit selects which can be profitable to convert back the branches.

PR: https://github.com/llvm/llvm-project/pull/115489

1 files changed, 21 insertions, 4 deletions

diff --git a/llvm/lib/CodeGen/SelectOptimize.cpp b/llvm/lib/CodeGen/SelectOptimize.cpp
index d480642..484705ea 100644
--- a/llvm/lib/CodeGen/SelectOptimize.cpp
+++ b/llvm/lib/CodeGen/SelectOptimize.cpp
@@ -489,8 +489,10 @@ static Value *getTrueOrFalseValue(
   }
 
   auto *BO = cast<BinaryOperator>(SI.getI());
-  assert(BO->getOpcode() == Instruction::Or &&
-         "Only currently handling Or instructions.");
+  assert((BO->getOpcode() == Instruction::Add ||
+          BO->getOpcode() == Instruction::Or ||
+          BO->getOpcode() == Instruction::Sub) &&
+         "Only currently handling Add, Or and Sub binary operators.");
 
   auto *CBO = BO->clone();
   auto CondIdx = SI.getConditionOpIndex();
@@ -786,8 +788,23 @@ void SelectOptimizeImpl::collectSelectGroups(BasicBlock &BB,
     // 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();
+      switch (I->getOpcode()) {
+      case Instruction::Add:
+      case Instruction::Sub: {
+        Value *X;
+        if (!((PatternMatch::match(I->getOperand(0),
+                                   m_OneUse(m_ZExt(m_Value(X)))) ||
+               PatternMatch::match(I->getOperand(1),
+                                   m_OneUse(m_ZExt(m_Value(X))))) &&
+              X->getType()->isIntegerTy(1)))
+          return SelectInfo.end();
+        break;
+      }
+      case Instruction::Or:
+        if (BO->getType()->isIntegerTy(1) || BO->getOpcode() != Instruction::Or)
+          return SelectInfo.end();
+        break;
+      }
 
       for (unsigned Idx = 0; Idx < 2; Idx++) {
         auto *Op = BO->getOperand(Idx);