[InstCombine] Don't fold select(C, Z, binop(select(C, X, Y), W)) -> select(C, Z, binop(Y, W)) if the binop is rem or div.

[craig.topper]

Summary

The select may have been preventing a division by zero so removing it might not be safe.

Fixes PR36362.

[spatel]

This is conservatively correct, so LGTM.

We could do better by using isSafeToSpeculativelyExecute(), but I think you'd have to use that after you've setOperand() and then switch it back if that fails. Or duplicate that logic with the proposed canMergeSelectThroughBinop().

The difference would show up in an example like this where we have a constant divisor, so we know there's no div-by-zero possibility.

define i32 @foo(i1 %a, i32 %b, i32 %c) {

%sel1 = select i1 %a, i32 42, i32 %b
%div = udiv i32 %c, %sel1
%sel2 = select i1 %a, i32 %div, i32 0
ret i32 %sel2

}

test/Transforms/InstCombine/pr36362.ll

9 ↗

(On Diff #134159)

Can minimize this to something like:

define i32 @foo(i1 %a, i32 %b, i32 %c) {
  %sel1 = select i1 %a, i32 %b, i32 -1
  %rem = srem i32 %c, %sel1
  %sel2 = select i1 %a, i32 %rem, i32 0
  ret i32 %sel2
}

llvm/trunk/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (FalseSI->getCondition()->getType() == CondVal->getType()) {
         Value *Or = Builder.CreateOr(CondVal, FalseSI->getCondition());
         SI.setOperand(0, Or);
         SI.setOperand(2, FalseSI->getFalseValue());
         return &SI;
       }
     }
   }
 
+  auto canMergeSelectThroughBinop = [](BinaryOperator *BO) {
+    // The select might be preventing a division by 0.
+    switch (BO->getOpcode()) {
+    default:
+      return true;
+    case Instruction::SRem:
+    case Instruction::URem:
+    case Instruction::SDiv:
+    case Instruction::UDiv:
+      return false;
+    }
+  };
+
   // Try to simplify a binop sandwiched between 2 selects with the same
   // condition.
   // select(C, binop(select(C, X, Y), W), Z) -> select(C, binop(X, W), Z)
   BinaryOperator *TrueBO;
-  if (match(TrueVal, m_OneUse(m_BinOp(TrueBO)))) {
+  if (match(TrueVal, m_OneUse(m_BinOp(TrueBO))) &&
+      canMergeSelectThroughBinop(TrueBO)) {
     if (auto *TrueBOSI = dyn_cast<SelectInst>(TrueBO->getOperand(0))) {
       if (TrueBOSI->getCondition() == CondVal) {
         TrueBO->setOperand(0, TrueBOSI->getTrueValue());
         Worklist.Add(TrueBO);
         return &SI;
       }
     }
     if (auto *TrueBOSI = dyn_cast<SelectInst>(TrueBO->getOperand(1))) {
       if (TrueBOSI->getCondition() == CondVal) {
         TrueBO->setOperand(1, TrueBOSI->getTrueValue());
         Worklist.Add(TrueBO);
         return &SI;
       }
     }
   }
 
   // select(C, Z, binop(select(C, X, Y), W)) -> select(C, Z, binop(Y, W))
   BinaryOperator *FalseBO;
-  if (match(FalseVal, m_OneUse(m_BinOp(FalseBO)))) {
+  if (match(FalseVal, m_OneUse(m_BinOp(FalseBO))) &&
+      canMergeSelectThroughBinop(FalseBO)) {
     if (auto *FalseBOSI = dyn_cast<SelectInst>(FalseBO->getOperand(0))) {
       if (FalseBOSI->getCondition() == CondVal) {
         FalseBO->setOperand(0, FalseBOSI->getFalseValue());
         Worklist.Add(FalseBO);
         return &SI;
       }
     }
     if (auto *FalseBOSI = dyn_cast<SelectInst>(FalseBO->getOperand(1))) {

llvm/trunk/test/Transforms/InstCombine/pr36362.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+;RUN: opt -instcombine -S %s | FileCheck %s
+
+; We shouldn't remove the select before the srem
+define i32 @foo(i1 %a, i32 %b, i32 %c) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[A:%.*]], i32 [[B:%.*]], i32 -1
+; CHECK-NEXT:    [[REM:%.*]] = srem i32 [[C:%.*]], [[SEL1]]
+; CHECK-NEXT:    [[SEL2:%.*]] = select i1 [[A]], i32 [[REM]], i32 0
+; CHECK-NEXT:    ret i32 [[SEL2]]
+;
+  %sel1 = select i1 %a, i32 %b, i32 -1
+  %rem = srem i32 %c, %sel1
+  %sel2 = select i1 %a, i32 %rem, i32 0
+  ret i32 %sel2
+}
+