[InstCombine] Fold (op x, ({z,s}ext (icmp eq x, C))) to select

goldsteinn · goldsteinn · commit 4601fa85576e · 2024-04-16T23:24:11.000-05:00
`(op x, ({z,s}ext (icmp eq x, C)))` is either `(op C, ({z,s}ext 1))`
or `(op x, 0)`.

If both possibilities simplify (i.e constant fold for the former and
either constant fold or converted to just `x` in the latter), fold to:

`(select (icmp eq x, C), (op C, ({z,s}ext 1)), (op x, 0)`.

Which is easier to analyze and should get roughly the same or better
codegen (in most cases something like `({z,s}ext (icmp))` lowers
simliarly to a `select`).
diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -1512,6 +1512,8 @@ bool llvm::canConstantFoldCallTo(const CallBase *Call, const Function *F) {
   case Intrinsic::usub_with_overflow:
   case Intrinsic::smul_with_overflow:
   case Intrinsic::umul_with_overflow:
+  case Intrinsic::sshl_sat:
+  case Intrinsic::ushl_sat:      
   case Intrinsic::sadd_sat:
   case Intrinsic::uadd_sat:
   case Intrinsic::ssub_sat:
@@ -2818,6 +2820,20 @@ static Constant *ConstantFoldIntrinsicCall2(Intrinsic::ID IntrinsicID, Type *Ty,
       };
       return ConstantStruct::get(cast<StructType>(Ty), Ops);
     }
+    case Intrinsic::sshl_sat:
+    case Intrinsic::ushl_sat:
+      // This is the same as for binary ops - poison propagates.
+      // TODO: Poison handling should be consolidated.
+      if (isa<PoisonValue>(Operands[0]) || isa<PoisonValue>(Operands[1]))
+        return PoisonValue::get(Ty);
+      if (!C0 && !C1)
+        return UndefValue::get(Ty);
+      if (!C0 || !C1)
+        return Constant::getNullValue(Ty);
+      if (IntrinsicID == Intrinsic::ushl_sat)
+        return ConstantInt::get(Ty, C0->ushl_sat(*C1));
+      else
+        return ConstantInt::get(Ty, C0->sshl_sat(*C1));
     case Intrinsic::uadd_sat:
     case Intrinsic::sadd_sat:
       // This is the same as for binary ops - poison propagates.
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -1458,6 +1458,9 @@ Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   // (A*B)+(A*C) -> A*(B+C) etc
   if (Value *V = foldUsingDistributiveLaws(I))
     return replaceInstUsesWith(I, V);
@@ -2076,6 +2079,9 @@ Instruction *InstCombinerImpl::visitSub(BinaryOperator &I) {
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
 
   // If this is a 'B = x-(-A)', change to B = x+A.
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -2275,6 +2275,9 @@ Instruction *InstCombinerImpl::visitAnd(BinaryOperator &I) {
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   // See if we can simplify any instructions used by the instruction whose sole
   // purpose is to compute bits we don't care about.
   if (SimplifyDemandedInstructionBits(I))
@@ -3444,6 +3447,9 @@ Instruction *InstCombinerImpl::visitOr(BinaryOperator &I) {
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   // See if we can simplify any instructions used by the instruction whose sole
   // purpose is to compute bits we don't care about.
   if (SimplifyDemandedInstructionBits(I))
@@ -4579,6 +4585,9 @@ Instruction *InstCombinerImpl::visitXor(BinaryOperator &I) {
   if (Instruction *NewXor = foldXorToXor(I, Builder))
     return NewXor;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   // (A&B)^(A&C) -> A&(B^C) etc
   if (Value *V = foldUsingDistributiveLaws(I))
     return replaceInstUsesWith(I, V);
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -1462,6 +1462,9 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(&CI);
   if (!II) return visitCallBase(CI);
 
+  if (Value *R = foldOpOfXWithXEqC(II, SQ.getWithInstruction(&CI)))
+    return replaceInstUsesWith(CI, R);
+
   // For atomic unordered mem intrinsics if len is not a positive or
   // not a multiple of element size then behavior is undefined.
   if (auto *AMI = dyn_cast<AtomicMemIntrinsic>(II))
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -754,6 +754,7 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
 
   Value *EvaluateInDifferentType(Value *V, Type *Ty, bool isSigned);
 
+  Value *foldOpOfXWithXEqC(Value *Op, const SimplifyQuery &SQ);
   bool tryToSinkInstruction(Instruction *I, BasicBlock *DestBlock);
   void tryToSinkInstructionDbgValues(
       Instruction *I, BasicBlock::iterator InsertPos, BasicBlock *SrcBlock,
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -204,6 +204,9 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   if (Value *V = foldUsingDistributiveLaws(I))
     return replaceInstUsesWith(I, V);
 
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -1020,6 +1020,9 @@ Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
   if (Instruction *V = dropRedundantMaskingOfLeftShiftInput(&I, Q, Builder))
     return V;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, Q))
+    return replaceInstUsesWith(I, R);
+
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   Type *Ty = I.getType();
   unsigned BitWidth = Ty->getScalarSizeInBits();
@@ -1256,6 +1259,9 @@ Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
   if (Instruction *R = commonShiftTransforms(I))
     return R;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   Type *Ty = I.getType();
   Value *X;
@@ -1591,6 +1597,9 @@ Instruction *InstCombinerImpl::visitAShr(BinaryOperator &I) {
   if (Instruction *R = commonShiftTransforms(I))
     return R;
 
+  if (Value *R = foldOpOfXWithXEqC(&I, SQ.getWithInstruction(&I)))
+    return replaceInstUsesWith(I, R);
+
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   Type *Ty = I.getType();
   unsigned BitWidth = Ty->getScalarSizeInBits();
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -4734,6 +4734,89 @@ void InstCombinerImpl::tryToSinkInstructionDbgValues(
   }
 }
 
+// If we have:
+//  `(op X, (zext/sext (icmp eq X, C)))`
+// We can transform it to:
+//  `(select (icmp eq X, C), (op C, (zext/sext 1)), (op X, 0))`
+// We do so if the `zext/sext` is one use and `(op X, 0)` simplifies.
+Value *InstCombinerImpl::foldOpOfXWithXEqC(Value *Op, const SimplifyQuery &SQ) {
+  Value *Cond;
+  Constant *C, *ExtC;
+
+  // match `(op X, (zext/sext (icmp eq X, C)))` and see if `(op X, 0)`
+  // simplifies.
+  // If we match and simplify, store the `icmp` in `Cond`, `(zext/sext C)` in
+  // `ExtC`.
+  auto MatchXWithXEqC = [&](Value *Op0, Value *Op1) -> Value * {
+    if (match(Op0, m_OneUse(m_ZExtOrSExt(m_Value(Cond))))) {
+      ICmpInst::Predicate Pred;
+      if (!match(Cond, m_ICmp(Pred, m_Specific(Op1), m_ImmConstant(C))) ||
+          Pred != ICmpInst::ICMP_EQ)
+        return nullptr;
+
+      ExtC = isa<SExtInst>(Op0) ? ConstantInt::getAllOnesValue(C->getType())
+                                : ConstantInt::get(C->getType(), 1);
+      return simplifyWithOpReplaced(Op, Op0,
+                                    Constant::getNullValue(Op1->getType()), SQ,
+                                    /*AllowRefinement=*/true);
+    }
+    return nullptr;
+  };
+
+  Value *SimpleOp = nullptr, *ConstOp = nullptr;
+  if (auto *BO = dyn_cast<BinaryOperator>(Op)) {
+    switch (BO->getOpcode()) {
+      // Potential TODO: For all of these, if Op1 is the compare, the compare
+      // must be true and we could replace Op0 with C (otherwise immediate UB).
+    case Instruction::UDiv:
+    case Instruction::SDiv:
+    case Instruction::URem:
+    case Instruction::SRem:
+      return nullptr;
+    default:
+      break;
+    }
+
+    // Try X is Op0
+    if ((SimpleOp = MatchXWithXEqC(BO->getOperand(0), BO->getOperand(1))))
+      ConstOp = Builder.CreateBinOp(BO->getOpcode(), ExtC, C);
+    // Try X is Op1
+    else if ((SimpleOp = MatchXWithXEqC(BO->getOperand(1), BO->getOperand(0))))
+      ConstOp = Builder.CreateBinOp(BO->getOpcode(), C, ExtC);
+  } else if (auto *II = dyn_cast<IntrinsicInst>(Op)) {
+    switch (II->getIntrinsicID()) {
+    default:
+      return nullptr;
+    case Intrinsic::sshl_sat:
+    case Intrinsic::ushl_sat:
+    case Intrinsic::umax:
+    case Intrinsic::umin:
+    case Intrinsic::smax:
+    case Intrinsic::smin:
+    case Intrinsic::uadd_sat:
+    case Intrinsic::usub_sat:
+    case Intrinsic::sadd_sat:
+    case Intrinsic::ssub_sat:
+      // Try X is Op0
+      if ((SimpleOp =
+               MatchXWithXEqC(II->getArgOperand(0), II->getArgOperand(1))))
+        ConstOp = Builder.CreateBinaryIntrinsic(II->getIntrinsicID(), ExtC, C);
+      // Try X is Op1
+      else if ((SimpleOp =
+                    MatchXWithXEqC(II->getArgOperand(1), II->getArgOperand(0))))
+        ConstOp = Builder.CreateBinaryIntrinsic(II->getIntrinsicID(), C, ExtC);
+      break;
+    }
+  }
+
+  assert((SimpleOp == nullptr) == (ConstOp == nullptr) &&
+         "Simplfied Op and Constant Op are de-synced!");
+  if (SimpleOp == nullptr)
+    return nullptr;
+
+  return Builder.CreateSelect(Cond, ConstOp, SimpleOp);
+}
+
 void InstCombinerImpl::tryToSinkInstructionDbgVariableRecords(
     Instruction *I, BasicBlock::iterator InsertPos, BasicBlock *SrcBlock,
     BasicBlock *DestBlock,
diff --git a/llvm/test/Transforms/InstCombine/apint-shift.ll b/llvm/test/Transforms/InstCombine/apint-shift.ll
@@ -564,14 +564,7 @@ define i40 @test26(i40 %A) {
 ; https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=9880
 define i177 @ossfuzz_9880(i177 %X) {
 ; CHECK-LABEL: @ossfuzz_9880(
-; CHECK-NEXT:    [[A:%.*]] = alloca i177, align 8
-; CHECK-NEXT:    [[L1:%.*]] = load i177, ptr [[A]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i177 [[L1]], -1
-; CHECK-NEXT:    [[B5_NEG:%.*]] = sext i1 [[TMP1]] to i177
-; CHECK-NEXT:    [[B14:%.*]] = add i177 [[L1]], [[B5_NEG]]
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i177 [[B14]], -1
-; CHECK-NEXT:    [[B1:%.*]] = zext i1 [[TMP2]] to i177
-; CHECK-NEXT:    ret i177 [[B1]]
+; CHECK-NEXT:    ret i177 0
 ;
   %A = alloca i177
   %L1 = load i177, ptr %A
diff --git a/llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll b/llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll
@@ -4,9 +4,7 @@
 declare void @use.i8(i8)
 define i8 @fold_add_zext_eq_0(i8 %x) {
 ; CHECK-LABEL: @fold_add_zext_eq_0(
-; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = zext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[X_EQ_EXT]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = call i8 @llvm.umax.i8(i8 [[X:%.*]], i8 1)
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 0
@@ -18,8 +16,7 @@ define i8 @fold_add_zext_eq_0(i8 %x) {
 define i8 @fold_add_sext_eq_0(i8 %x) {
 ; CHECK-LABEL: @fold_add_sext_eq_0(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = sext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[X_EQ_EXT]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[X_EQ]], i8 -1, i8 [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 0
@@ -73,8 +70,7 @@ define i8 @fold_mul_sext_eq_12_fail_multiuse(i8 %x) {
 define i8 @fold_shl_zext_eq_3_rhs(i8 %x) {
 ; CHECK-LABEL: @fold_shl_zext_eq_3_rhs(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = zext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = shl i8 [[X]], [[X_EQ_EXT]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[X_EQ]], i8 6, i8 [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 3
@@ -86,8 +82,7 @@ define i8 @fold_shl_zext_eq_3_rhs(i8 %x) {
 define i8 @fold_shl_zext_eq_3_lhs(i8 %x) {
 ; CHECK-LABEL: @fold_shl_zext_eq_3_lhs(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = zext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = shl nuw i8 [[X_EQ_EXT]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[X_EQ]], i8 8, i8 0
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 3
@@ -99,8 +94,7 @@ define i8 @fold_shl_zext_eq_3_lhs(i8 %x) {
 define <2 x i8> @fold_lshr_sext_eq_15_5_lhs(<2 x i8> %x) {
 ; CHECK-LABEL: @fold_lshr_sext_eq_15_5_lhs(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 15, i8 5>
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = sext <2 x i1> [[X_EQ]] to <2 x i8>
-; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i8> [[X_EQ_EXT]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[X_EQ]], <2 x i8> <i8 poison, i8 7>, <2 x i8> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %x_eq = icmp eq <2 x i8> %x, <i8 15, i8 5>
@@ -122,8 +116,7 @@ define <2 x i8> @fold_lshr_sext_eq_15_poison_rhs(<2 x i8> %x) {
 define i8 @fold_umax_zext_eq_9(i8 %x) {
 ; CHECK-LABEL: @fold_umax_zext_eq_9(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 9
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = sext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = call i8 @llvm.umax.i8(i8 [[X]], i8 [[X_EQ_EXT]])
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[X_EQ]], i8 -1, i8 [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 9
@@ -161,8 +154,7 @@ define i8 @fold_ushl_sat_zext_eq_3_lhs(i8 %x) {
 define i8 @fold_uadd_sat_zext_eq_3_rhs(i8 %x) {
 ; CHECK-LABEL: @fold_uadd_sat_zext_eq_3_rhs(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = zext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = call i8 @llvm.uadd.sat.i8(i8 [[X]], i8 [[X_EQ_EXT]])
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[X_EQ]], i8 4, i8 [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 3
@@ -187,8 +179,7 @@ define i8 @fold_ssub_sat_sext_eq_99_lhs_fail(i8 %x) {
 define i8 @fold_ssub_sat_zext_eq_99_rhs(i8 %x) {
 ; CHECK-LABEL: @fold_ssub_sat_zext_eq_99_rhs(
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 99
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = zext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = call i8 @llvm.ssub.sat.i8(i8 [[X]], i8 [[X_EQ_EXT]])
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[X_EQ]], i8 98, i8 [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 99
