[ValueTracking] Infer exact for udiv and sdiv

This infers the `exact` flag for division instructions based on
dominating conditions and assumes.

The `exact` flag is inferred from assumes for all targets but it won't
be inferred from dominating conditions for targets without a
single, unified instruction for division and remainder because
DivRemPairs rewrites the remainder instruction in terms of division.

Proof: https://alive2.llvm.org/ce/z/XtmMR6

Author: veera-sivarajan

llvm/include/llvm/Analysis/ValueTracking.h

@@ -122,6 +122,14 @@ bool isKnownToBeAPowerOfTwo(const Value *V, const DataLayout &DL,
 bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
                             const SimplifyQuery &Q);
 
+/// Return true if `X / Y` is known to be `exact` from dominating
+/// conditions.
+bool isKnownToBeAnExactDivFromConds(const Value *X, const Value *Y,
+                                    bool isSigned, const SimplifyQuery &Q);
+
+/// Return true if `V` is known to be equal to `0` from assumes.
+bool isKnownToBeZeroFromAssumes(const Value *V, const SimplifyQuery &Q);
+
 bool isOnlyUsedInZeroComparison(const Instruction *CxtI);
 
 bool isOnlyUsedInZeroEqualityComparison(const Instruction *CxtI);

llvm/include/llvm/Transforms/InstCombine/InstCombiner.h

@@ -447,6 +447,13 @@ class LLVM_LIBRARY_VISIBILITY InstCombiner {
                                         SQ.getWithInstruction(CxtI));
   }
 
+  bool isKnownToBeAnExactDivFromConds(const Value *X, const Value *Y,
+                                      bool isSigned,
+                                      const Instruction *CxtI = nullptr) {
+    return llvm::isKnownToBeAnExactDivFromConds(X, Y, isSigned,
+                                                SQ.getWithInstruction(CxtI));
+  }
+
   bool MaskedValueIsZero(const Value *V, const APInt &Mask, unsigned Depth = 0,
                          const Instruction *CxtI = nullptr) const {
     return llvm::MaskedValueIsZero(V, Mask, SQ.getWithInstruction(CxtI), Depth);

llvm/lib/Analysis/ValueTracking.cpp

@@ -2555,6 +2555,71 @@ bool llvm::isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
   }
 }
 
+bool llvm::isKnownToBeZeroFromAssumes(const Value *V, const SimplifyQuery &Q) {
+
+  if (Q.AC && Q.CxtI) {
+    for (auto &AssumeVH : Q.AC->assumptionsFor(V)) {
+      if (!AssumeVH)
+        continue;
+      CallInst *I = cast<CallInst>(AssumeVH);
+      const Value *Cond = I->getArgOperand(0);
+      if (match(Cond,
+                m_SpecificICmp(ICmpInst::ICMP_EQ, m_Specific(V), m_Zero())) &&
+          isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+        return true;
+      }
+    }
+  }
+
+  return false;
+}
+
+bool llvm::isKnownToBeAnExactDivFromConds(const Value *X, const Value *Y,
+                                          bool isSigned,
+                                          const SimplifyQuery &Q) {
+  if (Q.DC && Q.CxtI && Q.DT) {
+    auto MatchRem = [&](Value *DomCond, bool CondIsTrue) -> bool {
+      auto Pred = CondIsTrue ? ICmpInst::ICMP_EQ : ICmpInst::ICMP_NE;
+      if (isSigned)
+        return match(DomCond,
+                     m_SpecificICmp(Pred, m_SRem(m_Specific(X), m_Specific(Y)),
+                                    m_Zero()));
+      return match(
+          DomCond,
+          m_SpecificICmp(Pred, m_URem(m_Specific(X), m_Specific(Y)), m_Zero()));
+    };
+
+    auto *BB = Q.CxtI->getParent();
+    if (!Q.DT->isReachableFromEntry(BB))
+      return false;
+
+    auto *Node = Q.DT->getNode(BB);
+    while (Node->getIDom()) {
+      auto *Pred = Node->getIDom();
+      Node = Pred;
+
+      if (auto *BI = dyn_cast<BranchInst>(Pred->getBlock()->getTerminator());
+          BI && BI->isConditional()) {
+        Value *Cond = BI->getCondition();
+
+        BasicBlockEdge Edge0(BI->getParent(), BI->getSuccessor(0));
+        if (MatchRem(Cond, true) &&
+            Q.DT->dominates(Edge0, Q.CxtI->getParent())) {
+          return true;
+        }
+
+        BasicBlockEdge Edge1(BI->getParent(), BI->getSuccessor(1));
+        if (MatchRem(Cond, false) &&
+            Q.DT->dominates(Edge1, Q.CxtI->getParent())) {
+          return true;
+        }
+      }
+    }
+  }
+
+  return false;
+}
+
 /// Test whether a GEP's result is known to be non-null.
 ///
 /// Uses properties inherent in a GEP to try to determine whether it is known

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

@@ -1707,6 +1707,13 @@ Instruction *InstCombinerImpl::visitUDiv(BinaryOperator &I) {
     return replaceInstUsesWith(
         I, Builder.CreateLShr(Op0, Res, I.getName(), I.isExact()));
 
+  // Infer `exact` from dominating conditions.
+  if (!I.isExact() && isKnownToBeAnExactDivFromConds(Op0, Op1,
+                                                     /*isSigned=*/false, &I)) {
+    I.setIsExact();
+    return &I;
+  }
+
   return nullptr;
 }
 
@@ -1804,11 +1811,19 @@ Instruction *InstCombinerImpl::visitSDiv(BinaryOperator &I) {
   }
 
   KnownBits KnownDividend = computeKnownBits(Op0, 0, &I);
-  if (!I.isExact() &&
-      (match(Op1, m_Power2(Op1C)) || match(Op1, m_NegatedPower2(Op1C))) &&
-      KnownDividend.countMinTrailingZeros() >= Op1C->countr_zero()) {
-    I.setIsExact();
-    return &I;
+  if (!I.isExact()) {
+
+    if ((match(Op1, m_Power2(Op1C)) || match(Op1, m_NegatedPower2(Op1C))) &&
+        KnownDividend.countMinTrailingZeros() >= Op1C->countr_zero()) {
+      I.setIsExact();
+      return &I;
+    }
+
+    // Infer `exact` from dominating conditions.
+    if (isKnownToBeAnExactDivFromConds(Op0, Op1, /*isSigned=*/true, &I)) {
+      I.setIsExact();
+      return &I;
+    }
   }
 
   if (KnownDividend.isNonNegative()) {
@@ -1846,6 +1861,7 @@ Instruction *InstCombinerImpl::visitSDiv(BinaryOperator &I) {
     return SelectInst::Create(Cond, ConstantInt::get(Ty, 1),
                               ConstantInt::getAllOnesValue(Ty));
   }
+
   return nullptr;
 }
 

llvm/lib/Transforms/Scalar/DivRemPairs.cpp

@@ -15,6 +15,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -178,7 +179,7 @@ static DivRemWorklistTy getWorklist(Function &F) {
 /// Note: This transform could be an oddball enhancement to EarlyCSE, GVN, or
 /// SimplifyCFG, but it's split off on its own because it's different enough
 /// that it doesn't quite match the stated objectives of those passes.
-static bool optimizeDivRem(Function &F, const TargetTransformInfo &TTI,
+static bool optimizeDivRem(Function &F, TargetTransformInfo &TTI,
                            const DominatorTree &DT) {
   bool Changed = false;
 
@@ -196,6 +197,35 @@ static bool optimizeDivRem(Function &F, const TargetTransformInfo &TTI,
     auto &DivInst = E.DivInst;
     auto &RemInst = E.RemInst;
 
+    // Before any transformation, analyze whether
+    // RemInst is known to be zero from assumes. If so,
+    // set `exact` flag on DivInst.
+    //
+    // This kind of flag inference is usually handled by
+    // InstCombine but it will not work for this scenario
+    // because DivRemPairs is run towards the end of the
+    // optimization pipeline and removes all poison generating
+    // flags if the target does not have an unified instruction
+    // for both division and remainder.
+    //
+    // So, we do this flag inference here to ensure that flags
+    // inferred from assumes are not removed regardless of
+    // `HasDivRemOp`.
+    //
+    // Doing this in the tail end of the pipeline does not
+    // affect any follow-up optimizations because
+    // the `exact` flag is mostly used by the backend to
+    // generate better code.
+
+    const DataLayout &DL = DivInst->getDataLayout();
+    AssumptionCache AC(F, &TTI);
+    SimplifyQuery SQ(DL, &DT, &AC, DivInst);
+    bool RemInstIsZero = llvm::isKnownToBeZeroFromAssumes(RemInst, SQ);
+
+    if (RemInstIsZero) {
+      DivInst->setIsExact();
+    }
+
     const bool RemOriginallyWasInExpandedForm = E.isRemExpanded();
     (void)RemOriginallyWasInExpandedForm; // suppress unused variable warning
 
@@ -371,9 +401,11 @@ static bool optimizeDivRem(Function &F, const TargetTransformInfo &TTI,
       Sub->insertAfter(Mul->getIterator());
       Sub->setDebugLoc(RemInst->getDebugLoc());
 
-      // If DivInst has the exact flag, remove it. Otherwise this optimization
-      // may replace a well-defined value 'X % Y' with poison.
-      DivInst->dropPoisonGeneratingFlags();
+      // Remove the `exact` flag from DivInst if it is not inferred from
+      // assumes. Otherwise this optimization may replace a well-defined
+      // value 'X % Y' with poison.
+      if (!RemInstIsZero)
+        DivInst->dropPoisonGeneratingFlags();
 
       // If X can be undef, X should be frozen first.
       // For example, let's assume that Y = 1 & X = undef:

llvm/test/Transforms/DivRemPairs/PowerPC/div-rem-pairs-infer-exact.ll

@@ -10,7 +10,7 @@ define i8 @udiv_exact_assume(i8 %x, i8 %y) {
 ; CHECK-SAME: i8 [[X:%.*]], i8 [[Y:%.*]]) {
 ; CHECK-NEXT:    [[X_FROZEN:%.*]] = freeze i8 [[X]]
 ; CHECK-NEXT:    [[Y_FROZEN:%.*]] = freeze i8 [[Y]]
-; CHECK-NEXT:    [[DIV3:%.*]] = udiv i8 [[X_FROZEN]], [[Y_FROZEN]]
+; CHECK-NEXT:    [[DIV3:%.*]] = udiv exact i8 [[X_FROZEN]], [[Y_FROZEN]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = mul i8 [[DIV3]], [[Y_FROZEN]]
 ; CHECK-NEXT:    [[REM:%.*]] = sub i8 [[X_FROZEN]], [[TMP1]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[REM]], 0

llvm/test/Transforms/DivRemPairs/X86/div-rem-pairs-infer-exact.ll

@@ -10,7 +10,7 @@ define i8 @udiv_exact_assume(i8 %x, i8 %y) {
 ; CHECK-NEXT:    [[REM:%.*]] = urem i8 [[X]], [[Y]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[REM]], 0
 ; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[DIV3:%.*]] = udiv i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[DIV3:%.*]] = udiv exact i8 [[X]], [[Y]]
 ; CHECK-NEXT:    ret i8 [[DIV3]]
 ;
   %rem = urem i8 %x, %y
@@ -41,7 +41,7 @@ define i8 @infer_exact_from_dom_cond(i8 %X, i8 %Y) {
 ; CHECK-SAME: i8 [[X:%.*]], i8 [[Y:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    [[REM:%.*]] = srem i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv exact i8 [[X]], [[Y]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[REM]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label %[[IF_THEN:.*]], label %[[RETURN:.*]]
 ; CHECK:       [[IF_THEN]]:
@@ -69,7 +69,7 @@ define i8 @infer_exact_from_dom_cond_false_path(i8 %X, i8 %Y) {
 ; CHECK-SAME: i8 [[X:%.*]], i8 [[Y:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[REM:%.*]] = srem i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv exact i8 [[X]], [[Y]]
 ; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp eq i8 [[REM]], 0
 ; CHECK-NEXT:    br i1 [[CMP_NOT]], label %[[IF_ELSE:.*]], label %[[IF_THEN:.*]]
 ; CHECK:       [[IF_THEN]]: