[InstCombine] try to fold rem with constant dividend and select-of-constants divisor

[spatel]

Summary

We avoid this fold in the more general cases where we use FoldOpIntoSelect. That's because -- unlike most binary opcodes -- 'rem' can't usually be speculated with a variable divisor since it can have immediate UB. But in the case where both arms of the select are constants, we can safely evaluate both sides and eliminate 'rem' completely.

This should fix:
https://llvm.org/PR52102

I think we have the same optimization gap for 'div', but I need to add/run tests to make sure, so I'll do that as a follow-up if this is ok. We already handle the case where the dividend is a select and the divisor is constant - that's just above the diff in this patch.

[craig.topper]

Seems reasonable to me, Roman and Nikita are more active in InstCombine then I am lately so I'll let them look too.

[spatel]

Note: I drafted the equivalent patch for div opcodes, and it's exactly the same code/tests with s/rem/div/. No existing regression tests are affected.

[lebedev.ri]

Could you please add this test: https://godbolt.org/z/WvYxnejnc (it should not optimize to unreachable/undef/poison)

LGTM, but this reminds me of D71561 that got stuck (sorry, that was not my intention),
similar generalization could be applied here i believe.

[spatel]

In D115173#3177297, @lebedev.ri wrote:

Could you please add this test: https://godbolt.org/z/WvYxnejnc (it should not optimize to unreachable/undef/poison)

Sure - added "0 select arm" tests for all 4 opcodes:
51d3cb0ab15d

LGTM, but this reminds me of D71561 that got stuck (sorry, that was not my intention),
similar generalization could be applied here i believe.

I forgot about that patch...will try to revisit if I have time.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (Instruction *Op0I = dyn_cast<Instruction>(Op0)) {
       }
 
       // See if we can fold away this rem instruction.
       if (SimplifyDemandedInstructionBits(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)
+  if (match(Op0, m_ImmConstant()) &&
+      match(Op1, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()))) {
+    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op1)))
+      return R;
+  }
+
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::visitURem(BinaryOperator &I) {
   if (Value *V = SimplifyURemInst(I.getOperand(0), I.getOperand(1),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 

llvm/test/Transforms/InstCombine/rem.ll

 ;
   %sel = select i1 %cond, double %y, double 0.0
   %fmod = frem double %x, %sel
   ret double %fmod
 }
 
 define i32 @srem_constant_dividend_select_of_constants_divisor(i1 %b) {
 ; CHECK-LABEL: @srem_constant_dividend_select_of_constants_divisor(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 -3
-; CHECK-NEXT:    [[R:%.*]] = srem i32 42, [[S]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[B:%.*]], i32 6, i32 0
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 -3
   %r = srem i32 42, %s
   ret i32 %r
 }
 
+; TODO: srem should still be replaced by select.
+
 define i32 @srem_constant_dividend_select_of_constants_divisor_use(i1 %b) {
 ; CHECK-LABEL: @srem_constant_dividend_select_of_constants_divisor_use(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 -3
 ; CHECK-NEXT:    call void @use(i32 [[S]])
 ; CHECK-NEXT:    [[R:%.*]] = srem i32 42, [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 -3
   call void @use(i32 %s)
   %r = srem i32 42, %s
   ret i32 %r
 }
 
+; negative test - not safe to speculate rem with variable divisor
+
 define i32 @srem_constant_dividend_select_divisor1(i1 %b, i32 %x) {
 ; CHECK-LABEL: @srem_constant_dividend_select_divisor1(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 [[X:%.*]], i32 -3
 ; CHECK-NEXT:    [[R:%.*]] = srem i32 42, [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 %x, i32 -3
   %r = srem i32 42, %s
   ret i32 %r
 }
 
+; negative test - not safe to speculate rem with variable divisor
+
 define i32 @srem_constant_dividend_select_divisor2(i1 %b, i32 %x) {
 ; CHECK-LABEL: @srem_constant_dividend_select_divisor2(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 [[X:%.*]]
 ; CHECK-NEXT:    [[R:%.*]] = srem i32 42, [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 %x
   %r = srem i32 42, %s
   ret i32 %r
 }
 
 define <2 x i8> @srem_constant_dividend_select_of_constants_divisor_vec(i1 %b) {
 ; CHECK-LABEL: @srem_constant_dividend_select_of_constants_divisor_vec(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 4>
-; CHECK-NEXT:    [[R:%.*]] = srem <2 x i8> <i8 42, i8 -42>, [[S]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 6, i8 -2>, <2 x i8> <i8 2, i8 -2>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %s = select i1 %b, <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 4>
   %r = srem <2 x i8> <i8 42, i8 -42>, %s
   ret <2 x i8> %r
 }
 
+; Rem-by-0 element is immediate UB, so select is simplified.
+
 define <2 x i8> @srem_constant_dividend_select_of_constants_divisor_vec_ub1(i1 %b) {
 ; CHECK-LABEL: @srem_constant_dividend_select_of_constants_divisor_vec_ub1(
 ; CHECK-NEXT:    ret <2 x i8> <i8 2, i8 -2>
 ;
   %s = select i1 %b, <2 x i8> <i8 0, i8 -5>, <2 x i8> <i8 -4, i8 4>
   %r = srem <2 x i8> <i8 42, i8 -42>, %s
   ret <2 x i8> %r
 }
 
+; SMIN % -1 element is poison.
+
 define <2 x i8> @srem_constant_dividend_select_of_constants_divisor_vec_ub2(i1 %b) {
 ; CHECK-LABEL: @srem_constant_dividend_select_of_constants_divisor_vec_ub2(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 -1>
-; CHECK-NEXT:    [[R:%.*]] = srem <2 x i8> <i8 42, i8 -128>, [[S]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 6, i8 -3>, <2 x i8> <i8 2, i8 poison>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %s = select i1 %b, <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 -1>
   %r = srem <2 x i8> <i8 42, i8 -128>, %s
   ret <2 x i8> %r
 }
 
+; negative test - must have constant dividend
+
 define i32 @srem_select_of_constants_divisor(i1 %b, i32 %x) {
 ; CHECK-LABEL: @srem_select_of_constants_divisor(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 -3
 ; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X:%.*]], [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 -3
   %r = srem i32 %x, %s
   ret i32 %r
 }
 
 define i32 @urem_constant_dividend_select_of_constants_divisor(i1 %b) {
 ; CHECK-LABEL: @urem_constant_dividend_select_of_constants_divisor(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 -3
-; CHECK-NEXT:    [[R:%.*]] = urem i32 42, [[S]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[B:%.*]], i32 6, i32 42
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 -3
   %r = urem i32 42, %s
   ret i32 %r
 }
 
+; TODO: urem should still be replaced by select.
+
 define i32 @urem_constant_dividend_select_of_constants_divisor_use(i1 %b) {
 ; CHECK-LABEL: @urem_constant_dividend_select_of_constants_divisor_use(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 -3
 ; CHECK-NEXT:    call void @use(i32 [[S]])
 ; CHECK-NEXT:    [[R:%.*]] = urem i32 42, [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 -3
   call void @use(i32 %s)
   %r = urem i32 42, %s
   ret i32 %r
 }
 
+; negative test - not safe to speculate rem with variable divisor
+
 define i32 @urem_constant_dividend_select_divisor1(i1 %b, i32 %x) {
 ; CHECK-LABEL: @urem_constant_dividend_select_divisor1(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 [[X:%.*]], i32 -3
 ; CHECK-NEXT:    [[R:%.*]] = urem i32 42, [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 %x, i32 -3
   %r = urem i32 42, %s
   ret i32 %r
 }
 
+; negative test - not safe to speculate rem with variable divisor
+
 define i32 @urem_constant_dividend_select_divisor2(i1 %b, i32 %x) {
 ; CHECK-LABEL: @urem_constant_dividend_select_divisor2(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 [[X:%.*]]
 ; CHECK-NEXT:    [[R:%.*]] = urem i32 42, [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 %x
   %r = urem i32 42, %s
   ret i32 %r
 }
 
 define <2 x i8> @urem_constant_dividend_select_of_constants_divisor_vec(i1 %b) {
 ; CHECK-LABEL: @urem_constant_dividend_select_of_constants_divisor_vec(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 4>
-; CHECK-NEXT:    [[R:%.*]] = urem <2 x i8> <i8 42, i8 -42>, [[S]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 6, i8 -42>, <2 x i8> <i8 42, i8 2>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %s = select i1 %b, <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 4>
   %r = urem <2 x i8> <i8 42, i8 -42>, %s
   ret <2 x i8> %r
 }
 
+; Rem-by-0 element is immediate UB, so select is simplified.
+
 define <2 x i8> @urem_constant_dividend_select_of_constants_divisor_vec_ub1(i1 %b) {
 ; CHECK-LABEL: @urem_constant_dividend_select_of_constants_divisor_vec_ub1(
 ; CHECK-NEXT:    ret <2 x i8> <i8 42, i8 2>
 ;
   %s = select i1 %b, <2 x i8> <i8 0, i8 -5>, <2 x i8> <i8 -4, i8 4>
   %r = urem <2 x i8> <i8 42, i8 -42>, %s
   ret <2 x i8> %r
 }
 
+; There's no unsigned equivalent to "SMIN % -1", so this is just the usual constant folding.
+
 define <2 x i8> @urem_constant_dividend_select_of_constants_divisor_vec_ub2(i1 %b) {
 ; CHECK-LABEL: @urem_constant_dividend_select_of_constants_divisor_vec_ub2(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 -1>
-; CHECK-NEXT:    [[R:%.*]] = urem <2 x i8> <i8 42, i8 -128>, [[S]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[B:%.*]], <2 x i8> <i8 6, i8 -128>, <2 x i8> <i8 42, i8 -128>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %s = select i1 %b, <2 x i8> <i8 12, i8 -5>, <2 x i8> <i8 -4, i8 -1>
   %r = urem <2 x i8> <i8 42, i8 -128>, %s
   ret <2 x i8> %r
 }
 
+; negative test - must have constant dividend
+
 define i32 @urem_select_of_constants_divisor(i1 %b, i32 %x) {
 ; CHECK-LABEL: @urem_select_of_constants_divisor(
 ; CHECK-NEXT:    [[S:%.*]] = select i1 [[B:%.*]], i32 12, i32 -3
 ; CHECK-NEXT:    [[R:%.*]] = urem i32 [[X:%.*]], [[S]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = select i1 %b, i32 12, i32 -3
   %r = urem i32 %x, %s
   ret i32 %r
 }