[InstSimplify] fold exact divide to poison if it is known to not divide evenly

[spatel]

Summary

This is related to the discussion in D140665. I was looking over the demanded bits implementation in IR and noticed that we just bail out of a potential fold if a udiv is exact:
DemandedBits source link

Also, see tests added with 7f0c11509e8f.

Then, I saw that we could lose a fold to poison if we zap the exact with that transform, so this patch tries to catch that as a preliminary step.

Alive2 proofs:
https://alive2.llvm.org/ce/z/zCjKM7
https://alive2.llvm.org/ce/z/-tz_RK (trailing zeros must be "less-than")
https://alive2.llvm.org/ce/z/c9CMsJ (general proof and specific example)

[nikic]

llvm/lib/Analysis/InstructionSimplify.cpp

1158

This is not correct, Q.CxtI can be an unrelated instruction. You need to pass the exact flag down into the function.

[spatel]

llvm/lib/Analysis/InstructionSimplify.cpp

1158

Ah, I saw how some FMF were handled and followed that code pattern, but now I see how no-wrap is handled for other opcodes.

[spatel]

Patch updated:
Pass in a bool for exact-ness from the callers instead of using the context instruction in the SimplifyQuery. Test diffs are not changed.

[lebedev.ri]

Do we do this for right-shifts already?
If not, might be easier to start with that.

llvm/lib/Analysis/InstructionSimplify.cpp

1160

I think this is a better proof:

https://alive2.llvm.org/ce/z/zCjKM7
But not: https://alive2.llvm.org/ce/z/-tz_RK

5695–5696

/*IsExact=*/

llvm/test/Transforms/InstSimplify/div.ll

393–394

Is there a non-uniform test?
https://alive2.llvm.org/ce/z/Qkz5ac

[spatel]

In D140733#4018842, @lebedev.ri wrote:

Do we do this for right-shifts already?
If not, might be easier to start with that.

There is a variation on this theme for shifts (pasted code below), so I'll put a TODO on that.
I saw the missed division first, and I think we'd be more likely to spot bugs with a botched div if this manages to uncover any broken transforms.

// The low bit cannot be shifted out of an exact shift if it is set.
if (isExact) {
  KnownBits Op0Known =
      computeKnownBits(Op0, Q.DL, /*Depth=*/0, Q.AC, Q.CxtI, Q.DT);
  if (Op0Known.One[0])
    return Op0;
}

llvm/lib/Analysis/InstructionSimplify.cpp

1160

Do the extra function args just allow showing the TZ values/bounds more easily, or is there another benefit from writing it like that?

[spatel]

Patch updated:

1. Added non-splat vector tests.
2. Added code comments for bool function args.
3. Added TODO for exact shifts.

[lebedev.ri]

llvm/lib/Analysis/InstructionSimplify.cpp

1160

Do the extra function args just allow showing the TZ values/bounds more easily

Yes. That wasn't apparent to me in original test.

[lebedev.ri]

LG, thank you.

llvm/include/llvm/Analysis/InstructionSimplify.h

Value *simplifyFMAFMul(Value *LHS, Value *RHS, FastMathFlags FMF,
                        const SimplifyQuery &Q,
                        fp::ExceptionBehavior ExBehavior = fp::ebIgnore,
                        RoundingMode Rounding = RoundingMode::NearestTiesToEven);
 
 /// Given operands for a Mul, fold the result or return null.
 Value *simplifyMulInst(Value *LHS, Value *RHS, const SimplifyQuery &Q);
 
 /// Given operands for an SDiv, fold the result or return null.
-Value *simplifySDivInst(Value *LHS, Value *RHS, const SimplifyQuery &Q);
+Value *simplifySDivInst(Value *LHS, Value *RHS, bool IsExact,
+                        const SimplifyQuery &Q);
 
 /// Given operands for a UDiv, fold the result or return null.
-Value *simplifyUDivInst(Value *LHS, Value *RHS, const SimplifyQuery &Q);
+Value *simplifyUDivInst(Value *LHS, Value *RHS, bool IsExact,
+                        const SimplifyQuery &Q);
 
 /// Given operands for an FDiv, fold the result or return null.
 Value *
 simplifyFDivInst(Value *LHS, Value *RHS, FastMathFlags FMF,
                  const SimplifyQuery &Q,
                  fp::ExceptionBehavior ExBehavior = fp::ebIgnore,
                  RoundingMode Rounding = RoundingMode::NearestTiesToEven);
 

llvm/lib/Analysis/InstructionSimplify.cpp

static bool isDivZero(Value *X, Value *Y, const SimplifyQuery &Q,
 
   // Try again for any divisor:
   // Is the dividend unsigned less than the divisor?
   return isICmpTrue(ICmpInst::ICMP_ULT, X, Y, Q, MaxRecurse);
 }
 
 /// These are simplifications common to SDiv and UDiv.
 static Value *simplifyDiv(Instruction::BinaryOps Opcode, Value *Op0, Value *Op1,
-                          const SimplifyQuery &Q, unsigned MaxRecurse) {
+                          bool IsExact, const SimplifyQuery &Q,
+                          unsigned MaxRecurse) {
   if (Constant *C = foldOrCommuteConstant(Opcode, Op0, Op1, Q))
     return C;
 
   if (Value *V = simplifyDivRem(Opcode, Op0, Op1, Q, MaxRecurse))
     return V;
 
+  // If this is an exact divide by a constant, then the dividend (Op0) must have
+  // at least as many trailing zeros as the divisor to divide evenly. If it has
+  // less trailing zeros, then the result must be poison.
+  const APInt *DivC;
+  if (IsExact && match(Op1, m_APInt(DivC)) && DivC->countTrailingZeros()) {
+    KnownBits KnownOp0 = computeKnownBits(Op0, Q.DL, 0, Q.AC, Q.CxtI, Q.DT);
+    if (KnownOp0.countMaxTrailingZeros() < DivC->countTrailingZeros())
+      return PoisonValue::get(Op0->getType());
+  }
+
   bool IsSigned = Opcode == Instruction::SDiv;
 
   // (X rem Y) / Y -> 0
   if ((IsSigned && match(Op0, m_SRem(m_Value(), m_Specific(Op1)))) ||
       (!IsSigned && match(Op0, m_URem(m_Value(), m_Specific(Op1)))))
     return Constant::getNullValue(Op0->getType());
 
   // (X /u C1) /u C2 -> 0 if C1 * C2 overflow
static Value *simplifyRem(Instruction::BinaryOps Opcode, Value *Op0, Value *Op1,
   if (isDivZero(Op0, Op1, Q, MaxRecurse, Opcode == Instruction::SRem))
     return Op0;
 
   return nullptr;
 }
 
 /// Given operands for an SDiv, see if we can fold the result.
 /// If not, this returns null.
-static Value *simplifySDivInst(Value *Op0, Value *Op1, const SimplifyQuery &Q,
-                               unsigned MaxRecurse) {
+static Value *simplifySDivInst(Value *Op0, Value *Op1, bool IsExact,
+                               const SimplifyQuery &Q, unsigned MaxRecurse) {
   // If two operands are negated and no signed overflow, return -1.
   if (isKnownNegation(Op0, Op1, /*NeedNSW=*/true))
     return Constant::getAllOnesValue(Op0->getType());
 
-  return simplifyDiv(Instruction::SDiv, Op0, Op1, Q, MaxRecurse);
+  return simplifyDiv(Instruction::SDiv, Op0, Op1, IsExact, Q, MaxRecurse);
 }
 
-Value *llvm::simplifySDivInst(Value *Op0, Value *Op1, const SimplifyQuery &Q) {
-  return ::simplifySDivInst(Op0, Op1, Q, RecursionLimit);
+Value *llvm::simplifySDivInst(Value *Op0, Value *Op1, bool IsExact,
+                              const SimplifyQuery &Q) {
+  return ::simplifySDivInst(Op0, Op1, IsExact, Q, RecursionLimit);
 }
 
 /// Given operands for a UDiv, see if we can fold the result.
 /// If not, this returns null.
-static Value *simplifyUDivInst(Value *Op0, Value *Op1, const SimplifyQuery &Q,
-                               unsigned MaxRecurse) {
-  return simplifyDiv(Instruction::UDiv, Op0, Op1, Q, MaxRecurse);
+static Value *simplifyUDivInst(Value *Op0, Value *Op1, bool IsExact,
+                               const SimplifyQuery &Q, unsigned MaxRecurse) {
+  return simplifyDiv(Instruction::UDiv, Op0, Op1, IsExact, Q, MaxRecurse);
 }
 
-Value *llvm::simplifyUDivInst(Value *Op0, Value *Op1, const SimplifyQuery &Q) {
-  return ::simplifyUDivInst(Op0, Op1, Q, RecursionLimit);
+Value *llvm::simplifyUDivInst(Value *Op0, Value *Op1, bool IsExact,
+                              const SimplifyQuery &Q) {
+  return ::simplifyUDivInst(Op0, Op1, IsExact, Q, RecursionLimit);
 }
 
 /// Given operands for an SRem, see if we can fold the result.
 /// If not, this returns null.
 static Value *simplifySRemInst(Value *Op0, Value *Op1, const SimplifyQuery &Q,
                                unsigned MaxRecurse) {
   // If the divisor is 0, the result is undefined, so assume the divisor is -1.
   // srem Op0, (sext i1 X) --> srem Op0, -1 --> 0
if (Op0 == Op1)
     return Constant::getNullValue(Op0->getType());
 
   // undef >> X -> 0
   // undef >> X -> undef (if it's exact)
   if (Q.isUndefValue(Op0))
     return IsExact ? Op0 : Constant::getNullValue(Op0->getType());
 
   // The low bit cannot be shifted out of an exact shift if it is set.
+  // TODO: Generalize by counting trailing zeros (see fold for exact division).
   if (IsExact) {
     KnownBits Op0Known =
         computeKnownBits(Op0, Q.DL, /*Depth=*/0, Q.AC, Q.CxtI, Q.DT);
     if (Op0Known.One[0])
       return Op0;
   }
 
   return nullptr;
case Instruction::Add:
     return simplifyAddInst(LHS, RHS, /* IsNSW */ false, /* IsNUW */ false, Q,
                            MaxRecurse);
   case Instruction::Sub:
     return simplifySubInst(LHS, RHS,  /* IsNSW */ false, /* IsNUW */ false, Q,
                            MaxRecurse);
   case Instruction::Mul:
     return simplifyMulInst(LHS, RHS, Q, MaxRecurse);
   case Instruction::SDiv:
-    return simplifySDivInst(LHS, RHS, Q, MaxRecurse);
+    return simplifySDivInst(LHS, RHS, /* IsExact */ false, Q, MaxRecurse);
   case Instruction::UDiv:
-    return simplifyUDivInst(LHS, RHS, Q, MaxRecurse);
+    return simplifyUDivInst(LHS, RHS, /* IsExact */ false, Q, MaxRecurse);
   case Instruction::SRem:
     return simplifySRemInst(LHS, RHS, Q, MaxRecurse);
   case Instruction::URem:
     return simplifyURemInst(LHS, RHS, Q, MaxRecurse);
   case Instruction::Shl:
     return simplifyShlInst(LHS, RHS, /* IsNSW */ false, /* IsNUW */ false, Q,
                            MaxRecurse);
   case Instruction::LShr:
case Instruction::Sub:
     return simplifySubInst(NewOps[0], NewOps[1],
                            Q.IIQ.hasNoSignedWrap(cast<BinaryOperator>(I)),
                            Q.IIQ.hasNoUnsignedWrap(cast<BinaryOperator>(I)), Q);
   case Instruction::FMul:
     return simplifyFMulInst(NewOps[0], NewOps[1], I->getFastMathFlags(), Q);
   case Instruction::Mul:
     return simplifyMulInst(NewOps[0], NewOps[1], Q);
   case Instruction::SDiv:
-    return simplifySDivInst(NewOps[0], NewOps[1], Q);
+    return simplifySDivInst(NewOps[0], NewOps[1],
+                            Q.IIQ.isExact(cast<BinaryOperator>(I)), Q);
   case Instruction::UDiv:
-    return simplifyUDivInst(NewOps[0], NewOps[1], Q);
+    return simplifyUDivInst(NewOps[0], NewOps[1],
+                            Q.IIQ.isExact(cast<BinaryOperator>(I)), Q);
   case Instruction::FDiv:
     return simplifyFDivInst(NewOps[0], NewOps[1], I->getFastMathFlags(), Q);
   case Instruction::SRem:
     return simplifySRemInst(NewOps[0], NewOps[1], Q);
   case Instruction::URem:
     return simplifyURemInst(NewOps[0], NewOps[1], Q);
   case Instruction::FRem:
     return simplifyFRemInst(NewOps[0], NewOps[1], I->getFastMathFlags(), Q);

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (isa<Instruction>(D) && match(D, m_OneUse(m_ZExt(m_Value(X)))) &&
     // urem C, (zext X) --> zext (urem C', X)
     return new ZExtInst(Builder.CreateBinOp(Opcode, TruncC, X), Ty);
   }
 
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::visitUDiv(BinaryOperator &I) {
-  if (Value *V = simplifyUDivInst(I.getOperand(0), I.getOperand(1),
+  if (Value *V = simplifyUDivInst(I.getOperand(0), I.getOperand(1), I.isExact(),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   if (Instruction *X = foldVectorBinop(I))
     return X;
 
   // Handle the integer div common cases
   if (Instruction *Common = commonIDivTransforms(I))
if (takeLog2(Builder, Op1, /*Depth*/0, /*DoFold*/false)) {
     return replaceInstUsesWith(
         I, Builder.CreateLShr(Op0, Res, I.getName(), I.isExact()));
   }
 
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::visitSDiv(BinaryOperator &I) {
-  if (Value *V = simplifySDivInst(I.getOperand(0), I.getOperand(1),
+  if (Value *V = simplifySDivInst(I.getOperand(0), I.getOperand(1), I.isExact(),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   if (Instruction *X = foldVectorBinop(I))
     return X;
 
   // Handle the integer div common cases
   if (Instruction *Common = commonIDivTransforms(I))

llvm/test/Transforms/InstCombine/udiv-simplify.ll

 ; CHECK-NEXT:    [[U:%.*]] = udiv i8 [[A:%.*]], 12
 ; CHECK-NEXT:    ret i8 [[U]]
 ;
   %o = or i8 %a, 3
   %u = udiv i8 %o, 12
   ret i8 %u
 }
 
-; TODO: This can't divide evenly, so it is poison.
+; This can't divide evenly, so it is poison.
 
 define i8 @udiv_exact_demanded_low_bits_set(i8 %a) {
 ; CHECK-LABEL: @udiv_exact_demanded_low_bits_set(
-; CHECK-NEXT:    [[O:%.*]] = or i8 [[A:%.*]], 3
-; CHECK-NEXT:    [[U:%.*]] = udiv exact i8 [[O]], 12
-; CHECK-NEXT:    ret i8 [[U]]
+; CHECK-NEXT:    ret i8 poison
 ;
   %o = or i8 %a, 3
   %u = udiv exact i8 %o, 12
   ret i8 %u
 }
 
 ; All high bits are set, so this simplifies.
 

llvm/test/Transforms/InstSimplify/div.ll

 define i1 @const_urem_1() {
 ; CHECK-LABEL: @const_urem_1(
 ; CHECK-NEXT:    ret i1 false
 ;
   %rem = urem i1 ptrtoint (ptr @g to i1), ptrtoint (ptr @g2 to i1)
   ret i1 %rem
 }
 
+; Can't divide evenly, so create poison.
+
 define i8 @sdiv_exact_trailing_zeros(i8 %x) {
 ; CHECK-LABEL: @sdiv_exact_trailing_zeros(
-; CHECK-NEXT:    [[O:%.*]] = or i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[R:%.*]] = sdiv exact i8 [[O]], -42
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    ret i8 poison
 ;
   %o = or i8 %x, 1           ; odd number
   %r = sdiv exact i8 %o, -42 ; can't divide exactly
   ret i8 %r
 }
 
+; Negative test - could divide evenly.
+
 define i8 @sdiv_exact_trailing_zeros_eq(i8 %x) {
 ; CHECK-LABEL: @sdiv_exact_trailing_zeros_eq(
 ; CHECK-NEXT:    [[O:%.*]] = or i8 [[X:%.*]], 2
 ; CHECK-NEXT:    [[R:%.*]] = sdiv exact i8 [[O]], -42
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %o = or i8 %x, 2
   %r = sdiv exact i8 %o, -42
   ret i8 %r
 }
 
+; Negative test - must be exact div.
+
 define i8 @sdiv_trailing_zeros(i8 %x) {
 ; CHECK-LABEL: @sdiv_trailing_zeros(
 ; CHECK-NEXT:    [[O:%.*]] = or i8 [[X:%.*]], 1
 ; CHECK-NEXT:    [[R:%.*]] = sdiv i8 [[O]], -12
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %o = or i8 %x, 1
   %r = sdiv i8 %o, -12
   ret i8 %r
 }
 
+; TODO: Match non-splat vector constants.
+
+define <2 x i8> @sdiv_exact_trailing_zeros_nonuniform_vector(<2 x i8> %x) {
+; CHECK-LABEL: @sdiv_exact_trailing_zeros_nonuniform_vector(
+; CHECK-NEXT:    [[O:%.*]] = or <2 x i8> [[X:%.*]], <i8 3, i8 1>
+; CHECK-NEXT:    [[R:%.*]] = sdiv exact <2 x i8> [[O]], <i8 12, i8 2>
+; CHECK-NEXT:    ret <2 x i8> [[R]]
+;
+  %o = or <2 x i8> %x, <i8 3, i8 1>
+  %r = sdiv exact <2 x i8> %o, <i8 12, i8 2>
+  ret <2 x i8> %r
+}
+
+; Can't divide evenly, so create poison.
+
 define <2 x i8> @udiv_exact_trailing_zeros(<2 x i8> %x) {
 ; CHECK-LABEL: @udiv_exact_trailing_zeros(
-; CHECK-NEXT:    [[O:%.*]] = or <2 x i8> [[X:%.*]], <i8 3, i8 3>
-; CHECK-NEXT:    [[R:%.*]] = udiv exact <2 x i8> [[O]], <i8 12, i8 12>
-; CHECK-NEXT:    ret <2 x i8> [[R]]
+; CHECK-NEXT:    ret <2 x i8> poison
 ;
   %o = or <2 x i8> %x, <i8 3, i8 3>
   %r = udiv exact <2 x i8> %o, <i8 12, i8 12>  ; can't divide exactly
   ret <2 x i8> %r
 }
 
+; Negative test - could divide evenly.
+
 define <2 x i8> @udiv_exact_trailing_zeros_eq(<2 x i8> %x) {
 ; CHECK-LABEL: @udiv_exact_trailing_zeros_eq(
 ; CHECK-NEXT:    [[O:%.*]] = or <2 x i8> [[X:%.*]], <i8 28, i8 28>
 ; CHECK-NEXT:    [[R:%.*]] = udiv exact <2 x i8> [[O]], <i8 12, i8 12>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %o = or <2 x i8> %x, <i8 28, i8 28>
   %r = udiv exact <2 x i8> %o, <i8 12, i8 12>
   ret <2 x i8> %r
 }
 
+; Negative test - must be exact div.
+
 define i8 @udiv_trailing_zeros(i8 %x) {
 ; CHECK-LABEL: @udiv_trailing_zeros(
 ; CHECK-NEXT:    [[O:%.*]] = or i8 [[X:%.*]], 1
 ; CHECK-NEXT:    [[R:%.*]] = udiv i8 [[O]], 12
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %o = or i8 %x, 1
   %r = udiv i8 %o, 12
   ret i8 %r
 }
 
+; Negative test - only the first element is poison
+
+define <2 x i8> @udiv_exact_trailing_zeros_nonuniform_vector(<2 x i8> %x) {
+; CHECK-LABEL: @udiv_exact_trailing_zeros_nonuniform_vector(
+; CHECK-NEXT:    [[O:%.*]] = or <2 x i8> [[X:%.*]], <i8 3, i8 3>
+; CHECK-NEXT:    [[R:%.*]] = udiv exact <2 x i8> [[O]], <i8 12, i8 1>
+; CHECK-NEXT:    ret <2 x i8> [[R]]
+;
+  %o = or <2 x i8> %x, <i8 3, i8 3>
+  %r = udiv exact <2 x i8> %o, <i8 12, i8 1>
+  ret <2 x i8> %r
+}
+
 !0 = !{i32 0, i32 3}