[llvm] [ValueTracking] Implement udiv support for isKnownNonNullFromDominatingCondition (PR #67282)

[Dhruv Chawla via llvm-commits]

https://github.com/dc03 updated https://github.com/llvm/llvm-project/pull/67282

>From 10d2d543aeca53bdc87711740e584f679e81809e Mon Sep 17 00:00:00 2001
From: Dhruv Chawla <44582521+dc03 at users.noreply.github.com>
Date: Sun, 24 Sep 2023 23:11:32 +0530
Subject: [PATCH 1/2] [ValueTracking] Implement udiv support for
 isKnownNonNullFromDominatingCondition

The second operand of a udiv has to be non-null, as division by zero is
UB.

Proofs: https://alive2.llvm.org/ce/z/PAWdmP

Fixes https://github.com/llvm/llvm-project/issues/64240.
---
 llvm/lib/Analysis/ValueTracking.cpp           |  6 +++++
 .../ValueTracking/select-known-non-zero.ll    | 24 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 973943790c0aad2..f177321f28cfa2e 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2263,6 +2263,12 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         return true;
     }
 
+    if (const auto *Div = dyn_cast<BinaryOperator>(U);
+        Div && Div->getOpcode() == BinaryOperator::UDiv &&
+        Div->getOperand(1) == V &&
+        isValidAssumeForContext(cast<Instruction>(U), CtxI, DT))
+      return true;
+
     // Consider only compare instructions uniquely controlling a branch
     Value *RHS;
     CmpInst::Predicate Pred;
diff --git a/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll b/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
index 8b1d2fd0181d66c..a12925ede72ba42 100644
--- a/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
+++ b/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
@@ -393,3 +393,27 @@ define i1 @inv_select_v_sle_nonneg(i8 %v, i8 %C, i8 %y) {
   %r = icmp eq i8 %s, 0
   ret i1 %r
 }
+
+; Check dominance of the udiv over the icmp.
+define i64 @incorrect_safe_div(i64 %n, i64 %d) {
+; CHECK-LABEL: @incorrect_safe_div(
+; CHECK-NEXT:    [[TMP1:%.*]] = udiv i64 [[N:%.*]], [[D:%.*]]
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+  %1 = udiv i64 %n, %d
+  %2 = icmp eq i64 %d, 0
+  %3 = select i1 %2, i64 -1, i64 %1
+  ret i64 %3
+}
+
+; Check post-dominance of the udiv over the icmp.
+define i64 @incorrect_safe_div_post(i64 %n, i64 %d) {
+; CHECK-LABEL: @incorrect_safe_div_post(
+; CHECK-NEXT:    [[TMP1:%.*]] = udiv i64 [[N:%.*]], [[D:%.*]]
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+  %1 = icmp eq i64 %d, 0
+  %2 = udiv i64 %n, %d
+  %3 = select i1 %1, i64 -1, i64 %2
+  ret i64 %3
+}

>From 1ec021b2eea1156a04c2202d7ec281dc8ccd07f0 Mon Sep 17 00:00:00 2001
From: Dhruv Chawla <44582521+dc03 at users.noreply.github.com>
Date: Mon, 25 Sep 2023 12:18:49 +0530
Subject: [PATCH 2/2] Use pattern matching

---
 llvm/lib/Analysis/ValueTracking.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index f177321f28cfa2e..77a477a13df1595 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2263,9 +2263,7 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         return true;
     }
 
-    if (const auto *Div = dyn_cast<BinaryOperator>(U);
-        Div && Div->getOpcode() == BinaryOperator::UDiv &&
-        Div->getOperand(1) == V &&
+    if (match(U, m_UDiv(m_Value(), m_Specific(V))) &&
         isValidAssumeForContext(cast<Instruction>(U), CtxI, DT))
       return true;