[InstCombine] canonicalize urem as cmp+select

[Allen]

Summary

Fix https://github.com/llvm/llvm-project/issues/60546

[nikic]

Needs more tests for preconditions, e.g. what happens if the constant is not 1?

Note that there is also a more general form of this fold, see https://github.com/llvm/llvm-project/blob/697a162fa63df328ec9ca334636c5e85390b2bf0/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp#L806. The tricky bit here is that we would nee X < 2*Y, which is hard to establish as a precondition in this context.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1791

Why the nuw requirement? Seems to be fine without it? https://alive2.llvm.org/ce/z/uo7HMz

1792

Why is this specific to assumes? Can we do a generic simplifyICmp? (If we do so, Op0 needs to be frozen.)

[Allen]

Address comments:
a) check the assume with more generate API simplifyICmpInst
b) Add 3 test cases, including some Negative tests
c) Fix the checking and only do this when the return value is isOneValue()
d) Add the logic of frozen

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1791

Thanks, Both of them are added in test cases.

1792

Apply your comment, thanks

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1792

Would this be better? I always hate to see isa<> + cast<> pairs - it feels like we're duplicating work.

if (auto *CVal = dyn_cast_or_null<Constant>(Val)) {
  if (CVal->isOneValue()) {
  }
}

[Allen]

update to use dyn_cast_or_null according comment

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1792

Apply your review, thanks for your detailed example.

[nikic]

llvm/include/llvm/Analysis/InstructionSimplify.h

244

Why the new parameter? You can use SQ.withInstruction(I) to change the context instruction.

[Allen]

use SQ.getWithInstruction(&I) to to change the context instruction

[Allen]

llvm/include/llvm/Analysis/InstructionSimplify.h

244

Thanks your review, I update with SQ.getWithInstruction(&I).

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1793–1794

Use match(Val, m_One()) instead?

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

25

Drop noundef attributes

78

Missing tests where add constant is not 1.

It would be good to also have a test that is not based on assume, where %x < %n is known through some other way instead.

[Allen]

Address comments

• use match(Val, m_One()) to replace CVal->isOneValue()
• Drop the attribute noundef of function urem_assume_without_nuw, then it will see freeze
• Add a new negative case urem_assume_with_unexpected_const

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1789

This comment doesn't match the code.
Should it be:

// For "(X + 1) % Op1" and if (X u< Op1) => (X + 1) == Op1 ? 0 : X + 1

1794

This is called FrozenX, but it is freezing Op0. Should it freeze X directly as the name suggests?

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

78

Still no tests with no "assume" call. There should be a way to show the transform using math or bitwise logic ops and/or using a dominating condition?

[Allen]

Address comments:
a) Adding a new case urem_without_assume
b) Update the code comment to match the code
c) Rename the FreezeX to FreezeOp0 as the name suggests

[Allen]

Add a new case urem_with_dominating_condition

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1789

Thanks, apply it.

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

78

Add a case urem_with_dominating_condition, but now it is not supported, need a following patch to fix

[nikic]

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

6

Should also drop noundef here and on the return values.

106

%cmp is a dead instruction here. The request wasn't to just drop the assume, but to establish the %x < %n precondition without using one. It's a bit tricky to do this without also triggering other folds, but maybe this would work?

define i64 @test(i64 %arg) {
  %x = shl nuw i64 1, %arg
  %n = shl nuw i64 3, %arg
  %add = add i64 %x, 1
  %out = urem i64 %add, %n
  ret i64 %out
}

Unfortunately this example still doesn't show the need for freeze (https://alive2.llvm.org/ce/z/FJCdt4) ... I'm pretty sure it's needed in the general case, but I'm not sure how to come up with an example.

[spatel]

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

44

If it's not changing behavior of the optimization (and it is not in any cases here from what I see), use i8 or possibly even narrower type widths. This makes it more likely that Alive2 can run quickly on all tests to verify correctness.

[efriedma]

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

106

If I'm understanding your question correctly, the following demonstrates the need for freeze.

define i8 @src(i8 %arg, i8 %arg2) {
  %x = urem i8 %arg, %arg2
  %add = add i8 %x, 1
  %out = urem i8 %add, %arg2
  ret i8 %out
}

define i8 @tgt(i8 %arg, i8 %arg2) {
  %xx = urem i8 %arg, %arg2
  %x = freeze i8 %xx
  %add = add i8 %x, 1
  %cmp = icmp eq i8 %add, %arg2
  %out = select i1 %cmp, i8 0, i8 %add
  ret i8 %out
}

[nikic]

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

106

Thanks, that's exactly what I was looking for!

[Allen]

Address comments:
a) Drop all the attribute noundef on the argument and return value
b) Change all the value type to i8 as Alive2 can run quickly on all tests to verify correctness
c) Update case urem_without_assume

[Allen]

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

6

Thanks, delete all the noundef.

44

Thanks, I adapt all the value type into i8.

106

Thanks @nikic and @efriedma for your test case.

[nikic]

LGTM, but please pre-commit tests.

[Allen]

update as precommit test in D144372 according review.

llvm/include/llvm/Analysis/InstructionSimplify.h

 /// Given operands for an Or, fold the result or return null.
 Value *simplifyOrInst(Value *LHS, Value *RHS, const SimplifyQuery &Q);
 
 /// Given operands for an Xor, fold the result or return null.
 Value *simplifyXorInst(Value *LHS, Value *RHS, const SimplifyQuery &Q);
 
 /// Given operands for an ICmpInst, fold the result or return null.
 Value *simplifyICmpInst(unsigned Predicate, Value *LHS, Value *RHS,
-                        const SimplifyQuery &Q);
+                        const SimplifyQuery &Q,
+                        const Instruction *CxtI = nullptr);
 
 /// Given operands for an FCmpInst, fold the result or return null.
 Value *simplifyFCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
                         FastMathFlags FMF, const SimplifyQuery &Q);
 
 /// Given operands for a SelectInst, fold the result or return null.
 Value *simplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,
                           const SimplifyQuery &Q);

llvm/lib/Analysis/InstructionSimplify.cpp

                              const SimplifyQuery &, unsigned);
 static Value *simplifyBinOp(unsigned, Value *, Value *, const SimplifyQuery &,
                             unsigned);
 static Value *simplifyBinOp(unsigned, Value *, Value *, const FastMathFlags &,
                             const SimplifyQuery &, unsigned);
 static Value *simplifyCmpInst(unsigned, Value *, Value *, const SimplifyQuery &,
                               unsigned);
 static Value *simplifyICmpInst(unsigned Predicate, Value *LHS, Value *RHS,
-                               const SimplifyQuery &Q, unsigned MaxRecurse);
+                               const SimplifyQuery &Q, unsigned MaxRecurse,
+                               const Instruction *CxtI = nullptr);
 static Value *simplifyOrInst(Value *, Value *, const SimplifyQuery &, unsigned);
 static Value *simplifyXorInst(Value *, Value *, const SimplifyQuery &,
                               unsigned);
 static Value *simplifyCastInst(unsigned, Value *, Type *, const SimplifyQuery &,
                                unsigned);
 static Value *simplifyGEPInst(Type *, Value *, ArrayRef<Value *>, bool,
                               const SimplifyQuery &, unsigned);
 static Value *simplifySelectInst(Value *, Value *, Value *,

     // umax(A, B) <u umin(A, D) --> false
     if (Pred == CmpInst::ICMP_ULT)
       return getFalse(ITy);
   }
 
   return nullptr;
 }
 
-static Value *simplifyICmpWithDominatingAssume(CmpInst::Predicate Predicate,
-                                               Value *LHS, Value *RHS,
-                                               const SimplifyQuery &Q) {
+static Value *
+simplifyICmpWithDominatingAssume(CmpInst::Predicate Predicate, Value *LHS,
+                                 Value *RHS, const SimplifyQuery &Q,
+                                 const Instruction *CxtI) {
+  auto *Inst = (CxtI != nullptr) ? CxtI : Q.CxtI;
+
   // Gracefully handle instructions that have not been inserted yet.
-  if (!Q.AC || !Q.CxtI || !Q.CxtI->getParent())
+  if (!Q.AC || !Inst || !Inst->getParent())
     return nullptr;
 
   for (Value *AssumeBaseOp : {LHS, RHS}) {
     for (auto &AssumeVH : Q.AC->assumptionsFor(AssumeBaseOp)) {
       if (!AssumeVH)
         continue;
 
       CallInst *Assume = cast<CallInst>(AssumeVH);
       if (std::optional<bool> Imp = isImpliedCondition(
               Assume->getArgOperand(0), Predicate, LHS, RHS, Q.DL))
-        if (isValidAssumeForContext(Assume, Q.CxtI, Q.DT))
+        if (isValidAssumeForContext(Assume, Inst, Q.DT))
           return ConstantInt::get(getCompareTy(LHS), *Imp);
     }
   }
 
   return nullptr;
 }
 
 /// Given operands for an ICmpInst, see if we can fold the result.
 /// If not, this returns null.
 static Value *simplifyICmpInst(unsigned Predicate, Value *LHS, Value *RHS,
-                               const SimplifyQuery &Q, unsigned MaxRecurse) {
+                               const SimplifyQuery &Q, unsigned MaxRecurse,
+                               const Instruction *CxtI) {
   CmpInst::Predicate Pred = (CmpInst::Predicate)Predicate;
   assert(CmpInst::isIntPredicate(Pred) && "Not an integer compare!");
 
   if (Constant *CLHS = dyn_cast<Constant>(LHS)) {
     if (Constant *CRHS = dyn_cast<Constant>(RHS))
       return ConstantFoldCompareInstOperands(Pred, CLHS, CRHS, Q.DL, Q.TLI);
 
     // If we have a constant, make sure it is on the RHS.

   }
 
   if (Value *V = simplifyICmpWithBinOp(Pred, LHS, RHS, Q, MaxRecurse))
     return V;
 
   if (Value *V = simplifyICmpWithMinMax(Pred, LHS, RHS, Q, MaxRecurse))
     return V;
 
-  if (Value *V = simplifyICmpWithDominatingAssume(Pred, LHS, RHS, Q))
+  if (Value *V = simplifyICmpWithDominatingAssume(Pred, LHS, RHS, Q, CxtI))
     return V;
 
   // Simplify comparisons of related pointers using a powerful, recursive
   // GEP-walk when we have target data available..
   if (LHS->getType()->isPointerTy())
     if (auto *C = computePointerICmp(Pred, LHS, RHS, Q))
       return C;
   if (auto *CLHS = dyn_cast<PtrToIntOperator>(LHS))

   if (isa<PHINode>(LHS) || isa<PHINode>(RHS))
     if (Value *V = threadCmpOverPHI(Pred, LHS, RHS, Q, MaxRecurse))
       return V;
 
   return nullptr;
 }
 
 Value *llvm::simplifyICmpInst(unsigned Predicate, Value *LHS, Value *RHS,
-                              const SimplifyQuery &Q) {
-  return ::simplifyICmpInst(Predicate, LHS, RHS, Q, RecursionLimit);
+                              const SimplifyQuery &Q, const Instruction *CxtI) {
+  return ::simplifyICmpInst(Predicate, LHS, RHS, Q, RecursionLimit, CxtI);
 }
 
 /// Given operands for an FCmpInst, see if we can fold the result.
 /// If not, this returns null.
 static Value *simplifyFCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
                                FastMathFlags FMF, const SimplifyQuery &Q,
                                unsigned MaxRecurse) {
   CmpInst::Predicate Pred = (CmpInst::Predicate)Predicate;

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

   // max unsigned value. In that case, the remainder is 0:
   // urem Op0, (sext i1 X) --> (Op0 == -1) ? 0 : Op0
   Value *X;
   if (match(Op1, m_SExt(m_Value(X))) && X->getType()->isIntOrIntVectorTy(1)) {
     Value *Cmp = Builder.CreateICmpEQ(Op0, ConstantInt::getAllOnesValue(Ty));
     return SelectInst::Create(Cmp, ConstantInt::getNullValue(Ty), Op0);
   }
 
+  // If X / Y == 0, then (X + 1) % Y => (X + 1 == Y) ? 0 : X + 1 .
+  if (match(Op0, m_Add(m_Value(X), m_One()))) {
+    Value *Val = simplifyICmpInst(ICmpInst::ICMP_ULT, X, Op1, SQ, &I);
+    if (Val && isa<Constant>(Val))
+      if (cast<Constant>(Val)->isOneValue()) {
+        Value *FrozenX = Builder.CreateFreeze(Op0, Op0->getName() + ".frozen");
+        Value *Cmp = Builder.CreateICmpEQ(FrozenX, Op1);
+        return SelectInst::Create(Cmp, ConstantInt::getNullValue(Ty), FrozenX);
+      }
+  }
+
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::visitSRem(BinaryOperator &I) {
   if (Value *V = simplifySRemInst(I.getOperand(0), I.getOperand(1),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 

llvm/test/Transforms/InstCombine/urem-via-cmp-select.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+; https://alive2.llvm.org/ce/z/UNmz9j
+define noundef i64 @urem_assume(i64 noundef %x, i64 noundef %n) {
+; CHECK-LABEL: @urem_assume(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[X:%.*]], [[N:%.*]]
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[X]], 1
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 [[ADD]], [[N]]
+; CHECK-NEXT:    [[OUT:%.*]] = select i1 [[TMP0]], i64 0, i64 [[ADD]]
+; CHECK-NEXT:    ret i64 [[OUT]]
+;
+start:
+  %cmp = icmp ult i64 %x, %n
+  tail call void @llvm.assume(i1 %cmp)
+  %add = add nuw i64 %x, 1
+  %out = urem i64 %add, %n
+  ret i64 %out
+}
+
+; https://alive2.llvm.org/ce/z/uo7HMz
+define noundef i64 @urem_assume_without_nuw(i64 noundef %x, i64 noundef %n) {
+; CHECK-LABEL: @urem_assume_without_nuw(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[X:%.*]], [[N:%.*]]
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[X]], 1
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 [[ADD]], [[N]]
+; CHECK-NEXT:    [[OUT:%.*]] = select i1 [[TMP0]], i64 0, i64 [[ADD]]
+; CHECK-NEXT:    ret i64 [[OUT]]
+;
+start:
+  %cmp = icmp ult i64 %x, %n
+  tail call void @llvm.assume(i1 %cmp)
+  %add = add i64 %x, 1
+  %out = urem i64 %add, %n
+  ret i64 %out
+}
+
+define noundef i64 @urem_assume_eq(i64 noundef %x, i64 noundef %n) {
+; CHECK-LABEL: @urem_assume_eq(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[X:%.*]], [[N:%.*]]
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[X]], 1
+; CHECK-NEXT:    [[OUT:%.*]] = urem i64 [[ADD]], [[N]]
+; CHECK-NEXT:    ret i64 [[OUT]]
+;
+start:
+  %cmp = icmp eq i64 %x, %n
+  tail call void @llvm.assume(i1 %cmp)
+  %add = add i64 %x, 1
+  %out = urem i64 %add, %n
+  ret i64 %out
+}
+
+; Negative test: The assume is false
+define noundef i64 @urem_assume_ne(i64 noundef %x, i64 noundef %n) {
+; CHECK-LABEL: @urem_assume_ne(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[X:%.*]], [[N:%.*]]
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[ADD:%.*]] = add i64 [[X]], 1
+; CHECK-NEXT:    [[OUT:%.*]] = urem i64 [[ADD]], [[N]]
+; CHECK-NEXT:    ret i64 [[OUT]]
+;
+start:
+  %cmp = icmp ne i64 %x, %n
+  tail call void @llvm.assume(i1 %cmp)
+  %add = add i64 %x, 1
+  %out = urem i64 %add, %n
+  ret i64 %out
+}
+
+declare void @llvm.assume(i1 noundef)