[InstCombine] More aggressively try and fold irem/idiv/mul into selects.

[goldstein.w.n]

Summary

Don't only do constant cases, also try cases were we can't constant
fold both sides. In the same we can't constant fold both sides, only
do transform if the select is single use.

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[nikic]

This is not generally legal for div/rem, see https://alive2.llvm.org/ce/z/2iiR26 using one of your tests. This is only legal if the operation is speculatable.

[goldstein.w.n]

In D146350#4204295, @nikic wrote:

This is not generally legal for div/rem, see https://alive2.llvm.org/ce/z/2iiR26 using one of your tests. This is only legal if the operation is speculatable.

Is there a flag for speculatable?

[nikic]

In D146350#4222729, @goldstein.w.n wrote:

In D146350#4204295, @nikic wrote:

This is not generally legal for div/rem, see https://alive2.llvm.org/ce/z/2iiR26 using one of your tests. This is only legal if the operation is speculatable.

Is there a flag for speculatable?

There is isSafeToSpeculativelyExecute().

[nikic]

Err, I probably misunderstood the question: There are no div/rem variants that are always speculatable, we can only speculate them if we know that the divisor is non-zero (or the more complex conditions for the signed case are met).

[goldstein.w.n]

Only do rem/div if denominator known non zero

[goldstein.w.n]

In D146350#4232936, @nikic wrote:

Err, I probably misunderstood the question: There are no div/rem variants that are always speculatable, we can only speculate them if we know that the divisor is non-zero (or the more complex conditions for the signed case are met).

Done + test cases for it (*fail_no_speculation).

[goldstein.w.n]

Rebase

[goldstein.w.n]

ping.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

191

Comment that it returns MultiUse.

193

Matching against all-wildcards is a bit silly :)

1020

Please use isSafeToSpeculativelyExecute() here. This check is not sufficient for sdiv (due to INT_MIN / -1), and we don't want to duplicate logic.

1847

Same here.

llvm/test/Transforms/InstCombine/binop-select.ll

521

Also test the straightforward case where the mul has a constant operand? Or is that already handled somewhere else (and if so, can we remove it)?

[goldstein.w.n]

llvm/test/Transforms/InstCombine/binop-select.ll

521

hmm? Not sure what you mean by this comment.

[goldstein.w.n]

Use isSafeToSpeculate to check if allows on non-imm operands

[nikic]

llvm/test/Transforms/InstCombine/binop-select.ll

521

I mean something like (c ? x : C1) * C2, that isn't based on the select condition.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

206

Is this code in here so the case where the div is always folded away is always handled, as we're at worst going to relax to poison there and not introduce UB?

If so, I think it might be more elegant to sink this check into FoldOpIntoSelect and use the isSafeToSpeculativelyExecute variant that accepts operands, so we can check the operation that we're actually going to introduce (and thus also skip the check if we're not going to introduce one).

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

206

and use the isSafeToSpeculativelyExecute variant that accepts operands

Well, it looks like I imagined that variant... So ignore that bit (or leave a TODO for possible future improvement) and just take the bit about only checking it if we're going to introduce an op.

[xbolva00]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1844

Updatr comment?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

206

Without that variant we end up with regressions in:

LLVM :: Transforms/InstCombine/div.ll
LLVM :: Transforms/InstCombine/rem.ll

I'm happy to just add the variant and use that, but I'm not sure what to do with load/call instructions.
Maybe just return false in those cases as they aren't really needed here?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

206

There are regressions either way. Would prefer to keep it this way if thats okay with you.

llvm/test/Transforms/InstCombine/binop-select.ll

521

done although unaffected by this patch.

[goldstein.w.n]

Update comments

[xbolva00]

We generally fail to call 'foldBinOpIntoSelectOrPhi' in commonIDivTransforms and commonIRemTransforms

How hard would be to add support for PHIs in your shouldFoldOp helper? (Can be followup).

[goldstein.w.n]

In D146350#4313194, @xbolva00 wrote:

We generally fail to call 'foldBinOpIntoSelectOrPhi' in commonIDivTransforms and commonIRemTransforms

How hard would be to add support for PHIs in your shouldFoldOp helper? (Can be followup).

Well we match explicitly against m_Select(...) to check the all constant case so that wouldn't apply.
The speculative execution check can be reused, however.

Personally to add PHI support I'd create a new function "shouldFoldOpIntoX(...)" then have a variant
for select/phi and base the imm match on what X is.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

206

ping @nikic, okay with as is or okay with regressions or something else?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

206

ping2 @nikic.

[nikic]

I think this is basically okay if you drop the mul handling. As @xbolva00 pointed out (thanks, this is the bit I was missing here!) the mul case is already handled by foldBinOpIntoSelectOrPhi() and you're duplicating the transform for a special case now. If we want to extend handling to more cases, we should do so generically inside foldBinOpIntoSelectOrPhi(), not by repeating the transform with different arguments.

[goldstein.w.n]

In D146350#4359337, @nikic wrote:

I think this is basically okay if you drop the mul handling. As @xbolva00 pointed out (thanks, this is the bit I was missing here!) the mul case is already handled by foldBinOpIntoSelectOrPhi() and you're duplicating the transform for a special case now. If we want to extend handling to more cases, we should do so generically inside foldBinOpIntoSelectOrPhi(), not by repeating the transform with different arguments.

Ah, I had misunderstood his comment. I can move all the logics for div/mul/rem/etc.. into foldBinOpInstoSelectOrPhi. We actually call it for udiv as well (with really weak speculative checks).

[goldstein.w.n]

Move all new logic into foldBinOpIntoSelectOrPhi.

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase + remove issafetospeculativelyexecute patch dependency

[goldstein.w.n]

Use bespoke speculation tracking for div/rem

[goldstein.w.n]

(no rush) @nikic, I dropped the dependencies on the generic issafetospeculate... patch. Can take a look at this when you have the time.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

static Value *foldMulShl1(BinaryOperator &Mul, bool CommuteOperands,
   }
 
   return nullptr;
 }
 
 static Value *takeLog2(IRBuilderBase &Builder, Value *Op, unsigned Depth,
                        bool DoFold);
 
+static std::optional<bool> shouldFoldOpIntoSelect(BinaryOperator &I, Value *Op,
+                                                  Value *OpOther) {
+  if (match(Op, m_Select(m_Value(), m_Value(), m_Value())))
+    return match(OpOther, m_ImmConstant()) &&
+           match(Op, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()));
+
+  return std::nullopt;
+}
+
 Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   if (Value *V =
           simplifyMulInst(Op0, Op1, I.hasNoSignedWrap(), I.hasNoUnsignedWrap(),
                           SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   if (SimplifyAssociativeOrCommutative(I))
     return &I;
 
   if (Instruction *X = foldVectorBinop(I))
     return X;
 
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   if (takeLog2(Builder, Op1, /*Depth*/ 0, /*DoFold*/ false)) {
     Value *Res = takeLog2(Builder, Op1, /*Depth*/ 0, /*DoFold*/ true);
     BinaryOperator *Shl = BinaryOperator::CreateShl(Op0, Res);
     // We can only propegate nuw flag.
     Shl->setHasNoUnsignedWrap(HasNUW);
     return Shl;
   }
 
+  if (auto MultiUse = shouldFoldOpIntoSelect(I, Op0, Op1))
+    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op0), *MultiUse))
+      return R;
+
+  if (auto MultiUse = shouldFoldOpIntoSelect(I, Op1, Op0))
+    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op1), *MultiUse))
+      return R;
+
   bool Changed = false;
   if (!HasNSW && willNotOverflowSignedMul(Op0, Op1, I)) {
     Changed = true;
     I.setHasNoSignedWrap(true);
   }
 
   if (!HasNUW && willNotOverflowUnsignedMul(Op0, Op1, I)) {
     Changed = true;
Instruction *InstCombinerImpl::commonIDivTransforms(BinaryOperator &I) {
   // Handle cases involving: [su]div X, (select Cond, Y, Z)
   // This does not apply for fdiv.
   if (simplifyDivRemOfSelectWithZeroOp(I))
     return &I;
 
   // If the divisor is a select-of-constants, try to constant fold all div ops:
   // C / (select Cond, TrueC, FalseC) --> select Cond, (C / TrueC), (C / FalseC)
   // TODO: Adapt simplifyDivRemOfSelectWithZeroOp to allow this and other folds.
-  if (match(Op0, m_ImmConstant()) &&
-      match(Op1, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()))) {
-    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op1),
-                                          /*FoldWithMultiUse*/ true))
+  const SimplifyQuery Q = SQ.getWithInstruction(&I);
+  if (isKnownNonZero(Op1, DL, 0, Q.AC, Q.CxtI, Q.DT)) {
+    if (auto MultiUse = shouldFoldOpIntoSelect(I, Op0, Op1))
+      if (Instruction *R =
+              FoldOpIntoSelect(I, cast<SelectInst>(Op0), *MultiUse))
+        return R;
+
+    if (auto MultiUse = shouldFoldOpIntoSelect(I, Op1, Op0))
+      if (Instruction *R =
+              FoldOpIntoSelect(I, cast<SelectInst>(Op1), *MultiUse))
         return R;
   }
 
   const APInt *C2;
   if (match(Op1, m_APInt(C2))) {
     Value *X;
     const APInt *C1;
 
     // (X / C1) / C2  -> X / (C1*C2)
Instruction *InstCombinerImpl::commonIRemTransforms(BinaryOperator &I) {
   if (Value *V = simplifyValueKnownNonZero(I.getOperand(1), *this, I))
     return replaceOperand(I, 1, V);
 
   // Handle cases involving: rem X, (select Cond, Y, Z)
   if (simplifyDivRemOfSelectWithZeroOp(I))
     return &I;
 
   // If the divisor is a select-of-constants, try to constant fold all rem ops:
   // C % (select Cond, TrueC, FalseC) --> select Cond, (C % TrueC), (C % FalseC)
   // TODO: Adapt simplifyDivRemOfSelectWithZeroOp to allow this and other folds.
-  if (match(Op0, m_ImmConstant()) &&
-      match(Op1, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()))) {
-    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op1),
-                                          /*FoldWithMultiUse*/ true))
+  const SimplifyQuery Q = SQ.getWithInstruction(&I);
+  if (isKnownNonZero(Op1, DL, 0, Q.AC, Q.CxtI, Q.DT)) {
+    if (auto MultiUse = shouldFoldOpIntoSelect(I, Op0, Op1))
+      if (Instruction *R =
+              FoldOpIntoSelect(I, cast<SelectInst>(Op0), *MultiUse))
+        return R;
+
+    if (auto MultiUse = shouldFoldOpIntoSelect(I, Op1, Op0))
+      if (Instruction *R =
+              FoldOpIntoSelect(I, cast<SelectInst>(Op1), *MultiUse))
         return R;
   }
 
   if (isa<Constant>(Op1)) {
     if (Instruction *Op0I = dyn_cast<Instruction>(Op0)) {
       if (SelectInst *SI = dyn_cast<SelectInst>(Op0I)) {
         if (Instruction *R = FoldOpIntoSelect(I, SI))
           return R;
       } else if (auto *PN = dyn_cast<PHINode>(Op0I)) {

llvm/test/Transforms/InstCombine/binop-select.ll

;
   %s = select i1 %b, i32 2, i32 0
   call void @use(i32 %s)
   %r = ashr i32 -2, %s
   ret i32 %r
 }
 
 define <2 x i32> @test_udiv_to_const_shr(i1 %c, <2 x i32> %x, <2 x i32> %yy) {
 ; CHECK-LABEL: @test_udiv_to_const_shr(
-; CHECK-NEXT:    [[Y:%.*]] = shl nuw <2 x i32> <i32 1, i32 1>, [[YY:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C:%.*]], <2 x i32> <i32 7, i32 8>, <2 x i32> [[Y]]
-; CHECK-NEXT:    [[DIV:%.*]] = udiv <2 x i32> <i32 42, i32 42>, [[COND]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> <i32 42, i32 42>, [[YY:%.*]]
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C:%.*]], <2 x i32> <i32 6, i32 5>, <2 x i32> [[TMP1]]
 ; CHECK-NEXT:    ret <2 x i32> [[DIV]]
 ;
   %y = shl <2 x i32> <i32 1, i32 1>, %yy
   %cond = select i1 %c, <2 x i32> <i32 7, i32 8>, <2 x i32> %y
   %div = udiv <2 x i32> <i32 42, i32 42>, %cond
   ret <2 x i32> %div
 }
 
;
   %div = udiv i32 42, %cond
   call void @use(i32 %cond)
   ret i32 %div
 }
 
 define i32 @test_udiv_to_const_Cudiv(i32 %x) {
 ; CHECK-LABEL: @test_udiv_to_const_Cudiv(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 90
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 7, i32 19
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i32 [[X]], [[COND]]
+; CHECK-NEXT:    [[TMP1:%.*]] = udiv i32 [[X]], 19
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C]], i32 12, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %c = icmp eq i32 %x, 90
   %cond = select i1 %c, i32 7, i32 19
   %div = udiv i32 %x, %cond
   ret i32 %div
 }
 
;
   ret i32 %div
 }
 
 define i32 @test_udiv_partial_const(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test_udiv_partial_const(
 ; CHECK-NEXT:    [[YNZ:%.*]] = icmp ne i32 [[Y:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[YNZ]])
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 90
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 19, i32 [[Y]]
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i32 [[X]], [[COND]]
+; CHECK-NEXT:    [[TMP1:%.*]] = udiv i32 [[X]], [[Y]]
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C]], i32 4, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %ynz = icmp ne i32 %y, 0
   call void @llvm.assume(i1 %ynz)
   %c = icmp eq i32 %x, 90
   %cond = select i1 %c, i32 19, i32 %y
   %div = udiv i32 %x, %cond
   ret i32 %div
;
   %cond = select i1 %c, i32 7, i32 19
   %div = udiv i32 %x, %cond
   ret i32 %div
 }
 
 define <2 x i32> @test_sdiv_to_const_Csdiv(<2 x i32> %x) {
 ; CHECK-LABEL: @test_sdiv_to_const_Csdiv(
 ; CHECK-NEXT:    [[C_NOT:%.*]] = icmp eq <2 x i32> [[X:%.*]], <i32 90, i32 91>
-; CHECK-NEXT:    [[COND:%.*]] = select <2 x i1> [[C_NOT]], <2 x i32> <i32 16, i32 4>, <2 x i32> <i32 7, i32 9>
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv <2 x i32> [[X]], [[COND]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sdiv <2 x i32> [[X]], <i32 7, i32 9>
+; CHECK-NEXT:    [[DIV:%.*]] = select <2 x i1> [[C_NOT]], <2 x i32> <i32 5, i32 22>, <2 x i32> [[TMP1]]
 ; CHECK-NEXT:    ret <2 x i32> [[DIV]]
 ;
   %c = icmp ne <2 x i32> %x, <i32 90, i32 91>
   %cond = select <2 x i1> %c, <2 x i32> <i32 7, i32 9>, <2 x i32> <i32 16, i32 4>
   %div = sdiv <2 x i32> %x, %cond
   ret <2 x i32> %div
 }
 
 define i32 @test_srem_to_const_Csrem(i32 %x) {
 ; CHECK-LABEL: @test_srem_to_const_Csrem(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 24
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 7, i32 16
-; CHECK-NEXT:    [[DIV:%.*]] = srem i32 [[X]], [[COND]]
+; CHECK-NEXT:    [[TMP1:%.*]] = srem i32 [[X]], 16
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C]], i32 3, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %c = icmp eq i32 %x, 24
   %cond = select i1 %c, i32 7, i32 16
   %div = srem i32 %x, %cond
   ret i32  %div
 }
 
;
   %cond = select i1 %c, i32 7, i32 16
   %div = srem i32 %x, %cond
   call void @use(i32 %cond)
   ret i32 %div
 }
 
 define i32 @test_urem_to_const_and_ind_x(i32 %x, i32 %yy) {
 ; CHECK-LABEL: @test_urem_to_const_and_ind_x(
-; CHECK-NEXT:    [[Y:%.*]] = shl nuw i32 1, [[YY:%.*]]
 ; CHECK-NEXT:    [[C_NOT:%.*]] = icmp eq i32 [[X:%.*]], 24
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C_NOT]], i32 19, i32 [[Y]]
-; CHECK-NEXT:    [[DIV:%.*]] = urem i32 [[X]], [[COND]]
+; CHECK-NEXT:    [[NOTMASK:%.*]] = shl nsw i32 -1, [[YY:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i32 [[NOTMASK]], -1
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C_NOT]], i32 5, i32 [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %y = shl i32 1, %yy
   %c = icmp ne i32 %x, 24
   %cond = select i1 %c, i32 %y, i32 19
   %div = urem i32 %x, %cond
   ret i32 %div
 }
 
 define i32 @test_urem_to_const_and(i1 %c, i32 %yy) {
 ; CHECK-LABEL: @test_urem_to_const_and(
-; CHECK-NEXT:    [[Y:%.*]] = shl nuw i32 1, [[YY:%.*]]
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C:%.*]], i32 [[Y]], i32 19
-; CHECK-NEXT:    [[DIV:%.*]] = urem i32 44, [[COND]]
+; CHECK-NEXT:    [[NOTMASK:%.*]] = shl nsw i32 -1, [[YY:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[NOTMASK]], 44
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i32 [[TMP1]], 44
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C:%.*]], i32 [[TMP2]], i32 6
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %y = shl i32 1, %yy
   %cond = select i1 %c, i32 %y, i32 19
   %div = urem i32 44, %cond
   ret i32 %div
 }
 
 define i32 @test_urem_to_const_and_fail_no_speculation(i1 %c, i32 %yy) {
 ; CHECK-LABEL: @test_urem_to_const_and_fail_no_speculation(
 ; CHECK-NEXT:    [[Y:%.*]] = shl i32 2, [[YY:%.*]]
 ; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C:%.*]], i32 [[Y]], i32 19
 ; CHECK-NEXT:    [[DIV:%.*]] = urem i32 44, [[COND]]
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %y = shl i32 2, %yy
   %cond = select i1 %c, i32 %y, i32 19
   %div = urem i32 44, %cond
   ret i32 %div
 }
 
 define i32 @test_mul_to_const_Cmul(i32 %x) {
 ; CHECK-LABEL: @test_mul_to_const_Cmul(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 14
-; CHECK-NEXT:    [[DIV:%.*]] = mul i32 [[COND]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[X]], 14
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[C]], i32 549, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %c = icmp eq i32 %x, 61
   %cond = select i1 %c, i32 9, i32 14
   %div = mul i32 %x, %cond
   ret i32 %div
 }