[InstCombine] Fold `select` of `srem` and conditional add

[antoniofrighetto]

Summary

Simplify a pattern that may show up when computing the remainder of euclidean division.

Fixes: #64305

Proofs: https://alive2.llvm.org/ce/z/9_KG6c.

[antoniofrighetto]

Didn't manage to generalize the specific pattern out of it, not sure if this can be done here.

[goldstein.w.n]

A note about the proofs. You proved this when the remander is 8, not generically for any power of 2. Proofs with constants should generally look more like: https://alive2.llvm.org/ce/z/9_KG6c
i.e instead of declaring with the constant is, declare the constraints you have on it.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Can you remove the NSW from the comment / all the alive2 links. It makes it seem like the nsw is necessary for correctness.

2608

This also works for urem, unless that already gets properly folded, can you also support that here?

2609

Just use m_SpecificInt(I0) instead of I1 and drop this equality comparison.

Although really, this doesn't need to be an integer at all, it could be:

if (!(match(TrueVal, m_Add(m_Value(RemRes), m_Value(Remainder))) &&
        match(RemRes, m_SRem(m_Value(Op), m_Specific(m_Value(Remainder))) && FalseVal == RemRes &&
        && isKnownToBeAPowerOfTwo(Remainder, ...)))

2616

Can you move this check to before the other match. 1) think its clearer, 2) if you add the generic isKnownToBeAPowerOfTwo approach is good to exhaust all the cheap checks before the recursive call.

2618

getZExtValue is really okay here, what if its an i128?
You can build a ConstantInt directly from an APInt, so just *I0 should be okay here.

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

239

The tests should be precommit,.

They are also insufficient. Can you add some negative cases for: incorrect sign test, add with different constant, wrong operand positions in select.

Can you also add a test that use vecs, and add one where the datatype is i128 (as your current code seems to be buggy there).

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2608

If urem already gets properly folded, then if its in some clear place can you just update that codes to work for the srem case as well (instead of adding an entirely new sequence for srem).

[antoniofrighetto]

Thanks for reviewing it, and for the note! Updated proofs & changes, and rebased.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Dropped, thanks.

2608

urem is directly folded into an and upon entering visitURem as it can only take positive values. srem requires dedicated handling instead.

2616

Moved above, and moved from using m_Power2 to the generic isKnownToBeAPowerOfTwo.

2618

Right, getZExtValue does not really look good when handling i128. Moved to using getAllOnesValue, thanks.

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

239

Tests updated, I'll be pre-comitting them should they look OK.

[antoniofrighetto]

Separated tests & added diff.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

Please update you proofs to replace 8 with a variable C and use assume to provide its constraints (in your case pow2 or zero).

[antoniofrighetto]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

I had tried with something similar, but i128 seems to lead Alive2 to timeout, and there's no assume for vectors. Updated proofs with this, if it suffices.

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Use InstCombiner::isSignBitCheck so you can handle the case where the condition is inverted and the add and srem operands are swapped.

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Or does something already canonicalize to have the compare be ICMP SLT?

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Ok there is a canonicalization inside InstCombinerImpl::foldSelectInstWithICmp to prefer the SLT form. But it only triggers if the select is the only use of the condition.

This srem transform doesn't have any one use checks. So we might not want to rely on that canonicalization and should use InstCombiner::isSignBitCheck to handle both cases.

[antoniofrighetto]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

I suppose this would require to have something like:

ConstantInt *RHS;
bool TrueIfSigned = false;

if (!(match(CondVal, m_ICmp(Pred, m_Value(RemRes), m_ConstantInt(RHS))) &&
      IC.isSignBitCheck(Pred, RHS->getValue(), TrueIfSigned)))

And possibly also continue checking that RHS->getZExtValue() == 0. Wouldn't this make the code a bit slightly less readable? Considering the canonicalization, maybe we could use m_OneUse instead?

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Should be m_APInt instead of m_ConstantInt to support vector splats. See foldSelectToCopysign for similar. I don't think you need RHS->getZExtValue() == 0, since the SGT form uses -1 instead of 0.

2616

This uses isKnownToBeAPowerOfTwo, but the tests only check for constants. isKnownToBeAPowerOfTwo handles more than that.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

I had tried with something similar, but i128 seems to lead Alive2 to timeout, and there's no assume for vectors. Updated proofs with this, if it suffices.

Yeah, there is noneed to use i128 in the alive2 proofs, just having a test for it is fine.

2616

This uses isKnownToBeAPowerOfTwo, but the tests only check for constants. isKnownToBeAPowerOfTwo handles more than that.

+1

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2604

Personally think with a comment the readability will be fine. If its that big a concern we should add m_ICmpSignTest is its a fairly common match req.

[antoniofrighetto]

Updated & rebased!

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

Good point, updated tests with rem_euclid_non_const_pow2, thanks!

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

This needs to be TrueIfSigned ? (FalseVal == Rem) : (TrueVal == Rem)

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

Wouldn't we need this right after we find the compare?

if (!TrueIfSigned)
  std::swap(TrueVal, FalseVal)

[antoniofrighetto]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

As per rem_euclid_2 test, I think that if we are handling a SGT, we should check that the srem is in the true arm.

[antoniofrighetto]

Updated & rebased.

[antoniofrighetto]

@craig.topper, we are still dealing with a SGT as it has not been canonicalized into a SLT, right?

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

This line match(TrueVal, m_Add(m_Value(RemRes), m_Value(Remainder))) is also incorrect when TrueIfSigned is false. Which is why I suggested swapping the True and False value.

[antoniofrighetto]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2616

Right, ifdef'ing out the call to foldSelectInstWithICmp exhibits this clearly. Should be good now.

[antoniofrighetto]

Fixed support of inverse operands.

[antoniofrighetto]

Rebased & further commented.

[craig.topper]

LGTM

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2618

I think this fails to handle the case where the add operands are commuted? Can't happen for constant remainder, but I don't think anything would guarantee the remainder is on the RHS if it's variable.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (auto *I = dyn_cast<Instruction>(V))
       CandidateBlocks.insert(I->getParent());
 
   for (BasicBlock *BB : CandidateBlocks)
     if (auto *PN = foldSelectToPhiImpl(Sel, BB, DT, Builder))
       return PN;
   return nullptr;
 }
 
+/// Tries to reduce a pattern that arises when calculating the remainder of the
+/// Euclidean division. When the divisor is a power of two and is guaranteed not
+/// to be negative, a signed remainder can be folded with a bitwise and.
+///
+/// (x % n) < 0 ? (x % n) + n : (x % n)
+///    -> x & (n - 1)
+static Instruction *foldSelectWithSRem(SelectInst &SI, InstCombinerImpl &IC,
+                                       IRBuilderBase &Builder) {
+  Value *CondVal = SI.getCondition();
+  Value *TrueVal = SI.getTrueValue();
+  Value *FalseVal = SI.getFalseValue();
+
+  ICmpInst::Predicate Pred;
+  Value *Op, *RemRes, *Remainder;
+  const APInt *C;
+  bool TrueIfSigned = false;
+
+  if (!(match(CondVal, m_ICmp(Pred, m_Value(RemRes), m_APInt(C))) &&
+        IC.isSignBitCheck(Pred, *C, TrueIfSigned)))
+    return nullptr;
+
+  // If the sign bit is not set, we have a SGE/SGT comparison, and the operands
+  // of the select are inverted.
+  if (!TrueIfSigned)
+    std::swap(TrueVal, FalseVal);
+
+  // We are matching a quite specific pattern here:
+  // %rem = srem i32 %x, %n
+  // %cnd = icmp slt i32 %rem, 0
+  // %add = add i32 %rem, %n
+  // %sel = select i1 %cnd, i32 %add, i32 %rem
+  if (!(match(TrueVal, m_Add(m_Value(RemRes), m_Value(Remainder))) &&
+        match(RemRes, m_SRem(m_Value(Op), m_Specific(Remainder))) &&
+        IC.isKnownToBeAPowerOfTwo(Remainder, /*OrZero*/ true) &&
+        FalseVal == RemRes))
+    return nullptr;
+
+  Value *Add = Builder.CreateAdd(Remainder,
+                                 Constant::getAllOnesValue(RemRes->getType()));
+  return BinaryOperator::CreateAnd(Op, Add);
+}
+
 static Value *foldSelectWithFrozenICmp(SelectInst &Sel, InstCombiner::BuilderTy &Builder) {
   FreezeInst *FI = dyn_cast<FreezeInst>(Sel.getCondition());
   if (!FI)
     return nullptr;
 
   Value *Cond = FI->getOperand(0);
   Value *TrueVal = Sel.getTrueValue(), *FalseVal = Sel.getFalseValue();
 
Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   auto *FI = dyn_cast<Instruction>(FalseVal);
   if (TI && FI && TI->getOpcode() == FI->getOpcode())
     if (Instruction *IV = foldSelectOpOp(SI, TI, FI))
       return IV;
 
   if (Instruction *I = foldSelectExtConst(SI))
     return I;
 
+  if (Instruction *I = foldSelectWithSRem(SI, *this, Builder))
+    return I;
+
   // Fold (select C, (gep Ptr, Idx), Ptr) -> (gep Ptr, (select C, Idx, 0))
   // Fold (select C, Ptr, (gep Ptr, Idx)) -> (gep Ptr, (select C, 0, Idx))
   auto SelectGepWithBase = [&](GetElementPtrInst *Gep, Value *Base,
                                bool Swap) -> GetElementPtrInst * {
     Value *Ptr = Gep->getPointerOperand();
     if (Gep->getNumOperands() != 2 || Gep->getPointerOperand() != Base ||
         !Gep->hasOneUse())
       return nullptr;

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

;
   %r1 = urem i5 %x, %y
   %r2 = urem i5 %x, %z
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i32 @rem_euclid_1(i32 %0) {
 ; CHECK-LABEL: @rem_euclid_1(
-; CHECK-NEXT:    [[REM:%.*]] = srem i32 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[REM]], 0
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[REM]], 8
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[ADD]], i32 [[REM]]
+; CHECK-NEXT:    [[SEL:%.*]] = and i32 [[TMP0:%.*]], 7
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %rem = srem i32 %0, 8
   %cond = icmp slt i32 %rem, 0
   %add = add i32 %rem, 8
   %sel = select i1 %cond, i32 %add, i32 %rem
   ret i32 %sel
 }
 
 define i32 @rem_euclid_2(i32 %0) {
 ; CHECK-LABEL: @rem_euclid_2(
-; CHECK-NEXT:    [[REM:%.*]] = srem i32 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[REM]], 8
-; CHECK-NEXT:    [[COND1:%.*]] = icmp slt i32 [[REM]], 0
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND1]], i32 [[ADD]], i32 [[REM]]
+; CHECK-NEXT:    [[SEL:%.*]] = and i32 [[TMP0:%.*]], 7
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %rem = srem i32 %0, 8
   %cond = icmp sgt i32 %rem, -1
   %add = add i32 %rem, 8
   %sel = select i1 %cond, i32 %rem, i32 %add
   ret i32 %sel
 }
 
 define i32 @rem_euclid_wrong_sign_test(i32 %0) {
 ; CHECK-LABEL: @rem_euclid_wrong_sign_test(
 ; CHECK-NEXT:    [[REM:%.*]] = srem i32 [[TMP0:%.*]], 8
 ; CHECK-NEXT:    [[COND:%.*]] = icmp sgt i32 [[REM]], 0
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[REM]], 8
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[ADD]], i32 [[REM]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
;
   %cond = icmp slt i32 %rem, 0
   %add = add i32 %rem, 8
   %sel = select i1 %cond, i32 %rem, i32 %add
   ret i32 %sel
 }
 
 define <2 x i32> @rem_euclid_vec(<2 x i32> %0) {
 ; CHECK-LABEL: @rem_euclid_vec(
-; CHECK-NEXT:    [[REM:%.*]] = srem <2 x i32> [[TMP0:%.*]], <i32 8, i32 8>
-; CHECK-NEXT:    [[COND:%.*]] = icmp slt <2 x i32> [[REM]], zeroinitializer
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw <2 x i32> [[REM]], <i32 8, i32 8>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x i32> [[ADD]], <2 x i32> [[REM]]
+; CHECK-NEXT:    [[SEL:%.*]] = and <2 x i32> [[TMP0:%.*]], <i32 7, i32 7>
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %rem = srem <2 x i32> %0, <i32 8, i32 8>
   %cond = icmp slt <2 x i32> %rem, <i32 0, i32 0>
   %add = add <2 x i32> %rem, <i32 8, i32 8>
   %sel = select <2 x i1> %cond, <2 x i32> %add, <2 x i32> %rem
   ret <2 x i32> %sel
 }
 
 define i128 @rem_euclid_i128(i128 %0) {
 ; CHECK-LABEL: @rem_euclid_i128(
-; CHECK-NEXT:    [[REM:%.*]] = srem i128 [[TMP0:%.*]], 8
-; CHECK-NEXT:    [[COND:%.*]] = icmp slt i128 [[REM]], 0
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i128 [[REM]], 8
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i128 [[ADD]], i128 [[REM]]
+; CHECK-NEXT:    [[SEL:%.*]] = and i128 [[TMP0:%.*]], 7
 ; CHECK-NEXT:    ret i128 [[SEL]]
 ;
   %rem = srem i128 %0, 8
   %cond = icmp slt i128 %rem, 0
   %add = add i128 %rem, 8
   %sel = select i1 %cond, i128 %add, i128 %rem
   ret i128 %sel
 }
 
 define i8 @rem_euclid_non_const_pow2(i8 %0, i8 %1) {
 ; CHECK-LABEL: @rem_euclid_non_const_pow2(
-; CHECK-NEXT:    [[POW2:%.*]] = shl nuw i8 1, [[TMP0:%.*]]
-; CHECK-NEXT:    [[REM:%.*]] = srem i8 [[TMP1:%.*]], [[POW2]]
-; CHECK-NEXT:    [[COND:%.*]] = icmp slt i8 [[REM]], 0
-; CHECK-NEXT:    [[ADD:%.*]] = select i1 [[COND]], i8 [[POW2]], i8 0
-; CHECK-NEXT:    [[SEL:%.*]] = add i8 [[REM]], [[ADD]]
+; CHECK-NEXT:    [[NOTMASK:%.*]] = shl nsw i8 -1, [[TMP0:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i8 [[NOTMASK]], -1
+; CHECK-NEXT:    [[SEL:%.*]] = and i8 [[TMP3]], [[TMP1:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %pow2 = shl i8 1, %0
   %rem = srem i8 %1, %pow2
   %cond = icmp slt i8 %rem, 0
   %add = add i8 %rem, %pow2
   %sel = select i1 %cond, i8 %add, i8 %rem
   ret i8 %sel
 }