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

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

The same optimization for 'div' is planned as a follow-up patch.

Differential Revision: https://reviews.llvm.org/D115173

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 779d298..69a611e 100644 (file)
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -1461,6 +1461,15 @@ Instruction *InstCombinerImpl::commonIRemTransforms(BinaryOperator &I) {
   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)))
+      return R;
+  }
+
   if (isa<Constant>(Op1)) {
     if (Instruction *Op0I = dyn_cast<Instruction>(Op0)) {
       if (SelectInst *SI = dyn_cast<SelectInst>(Op0I)) {

diff --git a/llvm/test/Transforms/InstCombine/rem.ll b/llvm/test/Transforms/InstCombine/rem.ll
index cceb12a..adb1800 100644 (file)
--- a/llvm/test/Transforms/InstCombine/rem.ll
+++ b/llvm/test/Transforms/InstCombine/rem.ll
@@ -777,8 +777,7 @@ define double @PR34870(i1 %cond, double %x, double %y) {
 
 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
@@ -786,6 +785,8 @@ define i32 @srem_constant_dividend_select_of_constants_divisor(i1 %b) {
   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
@@ -799,6 +800,8 @@ define i32 @srem_constant_dividend_select_of_constants_divisor_use(i1 %b) {
   ret i32 %r
 }
 
+; Rem-by-0 is immediate UB, so select is simplified.
+
 define i32 @srem_constant_dividend_select_of_constants_divisor_0_arm(i1 %b) {
 ; CHECK-LABEL: @srem_constant_dividend_select_of_constants_divisor_0_arm(
 ; CHECK-NEXT:    ret i32 6
@@ -808,6 +811,8 @@ define i32 @srem_constant_dividend_select_of_constants_divisor_0_arm(i1 %b) {
   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
@@ -819,6 +824,8 @@ define i32 @srem_constant_dividend_select_divisor1(i1 %b, i32 %x) {
   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:%.*]]
@@ -832,8 +839,7 @@ define i32 @srem_constant_dividend_select_divisor2(i1 %b, i32 %x) {
 
 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>
@@ -841,6 +847,8 @@ define <2 x i8> @srem_constant_dividend_select_of_constants_divisor_vec(i1 %b) {
   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>
@@ -850,10 +858,11 @@ define <2 x i8> @srem_constant_dividend_select_of_constants_divisor_vec_ub1(i1 %
   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>
@@ -861,6 +870,8 @@ define <2 x i8> @srem_constant_dividend_select_of_constants_divisor_vec_ub2(i1 %
   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
@@ -874,8 +885,7 @@ define i32 @srem_select_of_constants_divisor(i1 %b, i32 %x) {
 
 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
@@ -883,6 +893,8 @@ define i32 @urem_constant_dividend_select_of_constants_divisor(i1 %b) {
   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
@@ -896,6 +908,8 @@ define i32 @urem_constant_dividend_select_of_constants_divisor_use(i1 %b) {
   ret i32 %r
 }
 
+; Rem-by-0 is immediate UB, so select is simplified.
+
 define i32 @urem_constant_dividend_select_of_constants_divisor_0_arm(i1 %b) {
 ; CHECK-LABEL: @urem_constant_dividend_select_of_constants_divisor_0_arm(
 ; CHECK-NEXT:    ret i32 6
@@ -905,6 +919,8 @@ define i32 @urem_constant_dividend_select_of_constants_divisor_0_arm(i1 %b) {
   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
@@ -916,6 +932,8 @@ define i32 @urem_constant_dividend_select_divisor1(i1 %b, i32 %x) {
   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:%.*]]
@@ -929,8 +947,7 @@ define i32 @urem_constant_dividend_select_divisor2(i1 %b, i32 %x) {
 
 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>
@@ -938,6 +955,8 @@ define <2 x i8> @urem_constant_dividend_select_of_constants_divisor_vec(i1 %b) {
   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>
@@ -947,10 +966,11 @@ define <2 x i8> @urem_constant_dividend_select_of_constants_divisor_vec_ub1(i1 %
   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>
@@ -958,6 +978,8 @@ define <2 x i8> @urem_constant_dividend_select_of_constants_divisor_vec_ub2(i1 %
   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