From 2b8d93f3c4d18ba6a50aceb007e77a0fa955ce56 Mon Sep 17 00:00:00 2001
From: Dhruv Chawla <44582521+dc03@users.noreply.github.com>
Date: Sun, 24 Sep 2023 23:11:32 +0530
Subject: [PATCH 1/4] [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 1e0281b3f1bd7..6a5fd2b3c40b6 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2264,6 +2264,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 8b1d2fd0181d6..a12925ede72ba 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 0d1a47fc99b1d0da08e5f922ed0c6a5460731255 Mon Sep 17 00:00:00 2001
From: Dhruv Chawla <44582521+dc03@users.noreply.github.com>
Date: Mon, 25 Sep 2023 12:18:49 +0530
Subject: [PATCH 2/4] 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 6a5fd2b3c40b6..d3a9d66243a79 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2264,9 +2264,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;
 

From ee12907f34efdf4051d86252d636bbf26a70e9c2 Mon Sep 17 00:00:00 2001
From: Dhruv Chawla <44582521+dc03@users.noreply.github.com>
Date: Mon, 25 Sep 2023 17:24:24 +0530
Subject: [PATCH 3/4] Add sdiv support + Add more tests, clarifying terminology

---
 llvm/lib/Analysis/ValueTracking.cpp           |  3 +-
 .../ValueTracking/select-known-non-zero.ll    | 44 ++++++++++++++++---
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index d3a9d66243a79..736472d151ceb 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2264,7 +2264,8 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         return true;
     }
 
-    if (match(U, m_UDiv(m_Value(), m_Specific(V))) &&
+    if ((match(U, m_UDiv(m_Value(), m_Specific(V))) ||
+         match(U, m_SDiv(m_Value(), m_Specific(V)))) &&
         isValidAssumeForContext(cast<Instruction>(U), CtxI, DT))
       return true;
 
diff --git a/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll b/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
index a12925ede72ba..1dc88412041d3 100644
--- a/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
+++ b/llvm/test/Analysis/ValueTracking/select-known-non-zero.ll
@@ -394,9 +394,9 @@ define i1 @inv_select_v_sle_nonneg(i8 %v, i8 %C, i8 %y) {
   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 udiv/sdiv occuring before icmp.
+define i64 @incorrect_safe_div_1(i64 %n, i64 %d) {
+; CHECK-LABEL: @incorrect_safe_div_1(
 ; CHECK-NEXT:    [[TMP1:%.*]] = udiv i64 [[N:%.*]], [[D:%.*]]
 ; CHECK-NEXT:    ret i64 [[TMP1]]
 ;
@@ -406,14 +406,44 @@ define i64 @incorrect_safe_div(i64 %n, i64 %d) {
   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 icmp occuring before udiv/sdiv.
+define i64 @incorrect_safe_div_2(i64 %n, i64 %d) {
+; CHECK-LABEL: @incorrect_safe_div_2(
+; CHECK-NEXT:    [[TMP1:%.*]] = sdiv i64 [[N:%.*]], [[D:%.*]]
 ; CHECK-NEXT:    ret i64 [[TMP1]]
 ;
   %1 = icmp eq i64 %d, 0
+  %2 = sdiv i64 %n, %d
+  %3 = select i1 %1, i64 -1, i64 %2
+  ret i64 %3
+}
+
+define i64 @incorrect_safe_div_call_1(i64 %n, i64 %d) {
+; CHECK-LABEL: @incorrect_safe_div_call_1(
+; CHECK-NEXT:    [[TMP1:%.*]] = sdiv i64 [[N:%.*]], [[D:%.*]]
+; CHECK-NEXT:    tail call void @use(i64 [[D]])
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+  %1 = sdiv i64 %n, %d
+  tail call void @use(i64 %d)
+  %2 = icmp eq i64 %d, 0
+  %3 = select i1 %2, i64 -1, i64 %1
+  ret i64 %3
+}
+
+define i64 @incorrect_safe_div_call_2(i64 %n, i64 %d) {
+; CHECK-LABEL: @incorrect_safe_div_call_2(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[D:%.*]], 0
+; CHECK-NEXT:    tail call void @use(i64 [[D]])
+; CHECK-NEXT:    [[TMP2:%.*]] = udiv i64 [[N:%.*]], [[D]]
+; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 -1, i64 [[TMP2]]
+; CHECK-NEXT:    ret i64 [[TMP3]]
+;
+  %1 = icmp eq i64 %d, 0
+  tail call void @use(i64 %d)
   %2 = udiv i64 %n, %d
   %3 = select i1 %1, i64 -1, i64 %2
   ret i64 %3
 }
+
+declare void @use(i64)

From 9ffa883e6fca7557a6f4bc03f85bb3bd79344a81 Mon Sep 17 00:00:00 2001
From: Dhruv Chawla <44582521+dc03@users.noreply.github.com>
Date: Thu, 28 Sep 2023 13:44:29 +0530
Subject: [PATCH 4/4] Replace UDiv/SDiv matchers by IDiv

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

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 736472d151ceb..8824a05e3aa6c 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2264,8 +2264,7 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         return true;
     }
 
-    if ((match(U, m_UDiv(m_Value(), m_Specific(V))) ||
-         match(U, m_SDiv(m_Value(), m_Specific(V)))) &&
+    if (match(U, m_IDiv(m_Value(), m_Specific(V))) &&
         isValidAssumeForContext(cast<Instruction>(U), CtxI, DT))
       return true;