[llvm] [SelectionDAG] Add more cases for UDIV and SDIV (PR #86452)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/86452

>From 69339fc078b5f1b099a9a649a68d256f78152406 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sun, 24 Mar 2024 14:09:28 -0400
Subject: [PATCH] [SelectionDAG] Add more cases for UDIV and SDIV

Ported from ValueTracking.
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  23 +++-
 .../Analysis/ValueTracking/known-non-zero.ll  | 114 ++++++++++++++++++
 2 files changed, 134 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index df30148b78b65..630a8a8139abb 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5857,14 +5857,31 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
       return true;
     break;
   }
-  case ISD::UDIV:
-  case ISD::SDIV:
+  case ISD::UDIV: {
     // div exact can only produce a zero if the dividend is zero.
-    // TODO: For udiv this is also true if Op1 u<= Op0
     if (Op->getFlags().hasExact())
       return isKnownNeverZero(Op.getOperand(0), Depth + 1);
+
+    // If Op0 >= Op1, then the result is at least 1, and therefore not 0.
+    KnownBits Op0 = computeKnownBits(Op.getOperand(0), Depth + 1);
+    KnownBits Op1 = computeKnownBits(Op.getOperand(1), Depth + 1);
+    if (Op0.isStrictlyPositive() && Op1.isStrictlyPositive() &&
+        KnownBits::uge(Op0, Op1).value_or(false))
+      return true;
     break;
+  }
+  case ISD::SDIV: {
+    // div exact can only produce a zero if the dividend is zero.
+    if (Op->getFlags().hasExact())
+      return isKnownNeverZero(Op.getOperand(0), Depth + 1);
 
+    KnownBits Op0 = computeKnownBits(Op.getOperand(0), Depth + 1);
+    KnownBits Op1 = computeKnownBits(Op.getOperand(1), Depth + 1);
+    if (Op0.isStrictlyPositive() && Op1.isStrictlyPositive() &&
+        KnownBits::uge(Op0, Op1).value_or(false))
+      return true;
+    break;
+  }
   case ISD::ADD:
     if (Op->getFlags().hasNoUnsignedWrap())
       if (isKnownNeverZero(Op.getOperand(1), Depth + 1) ||
diff --git a/llvm/test/Analysis/ValueTracking/known-non-zero.ll b/llvm/test/Analysis/ValueTracking/known-non-zero.ll
index f02ff583e5cbb..d26cc78a38043 100644
--- a/llvm/test/Analysis/ValueTracking/known-non-zero.ll
+++ b/llvm/test/Analysis/ValueTracking/known-non-zero.ll
@@ -1547,3 +1547,117 @@ define i1 @vec_reverse_non_zero_demanded_fail(<4 x i8> %xx) {
 }
 
 declare i32 @llvm.experimental.get.vector.length.i32(i32, i32, i1)
+
+define i1 @sdiv_by_1_never_zero(i8 %x) {
+; CHECK-LABEL: @sdiv_by_1_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %nonzero = or i8 %x, 1 ; Make sure x is non-zero
+  %div = sdiv i8 %nonzero, 1
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @sdiv_by_neg1_never_zero(i8 %x) {
+; CHECK-LABEL: @sdiv_by_neg1_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %nonzero = or i8 %x, 1 ; Make sure x is non-zero
+  %div = sdiv i8 %nonzero, -1
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @udiv_by_1_never_zero(i8 %x) {
+; CHECK-LABEL: @udiv_by_1_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %nonzero = or i8 %x, 1 ; Make sure x is non-zero
+  %div = udiv i8 %nonzero, 1
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @sdiv_odd_by_pow2_never_zero(i8 %x) {
+; CHECK-LABEL: @sdiv_odd_by_pow2_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %odd = or i8 %x, 1 ; Make sure x is odd
+  %div = sdiv i8 %odd, 4
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @udiv_odd_by_pow2_never_zero(i8 %x) {
+; CHECK-LABEL: @udiv_odd_by_pow2_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %odd = or i8 %x, 1 ; Make sure x is odd
+  %div = udiv i8 %odd, 4
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+
+define i1 @sdiv_odd_by_neg_pow2_never_zero(i8 %x) {
+; CHECK-LABEL: @sdiv_odd_by_neg_pow2_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %odd = or i8 %x, 1 ; Make sure x is odd
+  %div = sdiv i8 %odd, -4
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @udiv_dividend_ge_divisor_never_zero(i8 %x, i8 %y) {
+; CHECK-LABEL: @udiv_dividend_ge_divisor_never_zero(
+; CHECK-NEXT:    [[NUM:%.*]] = or i8 [[X:%.*]], 16
+; CHECK-NEXT:    [[DENUM:%.*]] = and i8 [[Y:%.*]], 15
+; CHECK-NEXT:    [[DIV:%.*]] = udiv i8 [[NUM]], [[DENUM]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[DIV]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %num = or i8 %x, 16   ; Ensure dividend is at least 16
+  %denum = and i8 %y, 15 ; Ensure divisor is at most 15
+  %div = udiv i8 %num, %denum
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @sdiv_abs_dividend_ge_abs_divisor_never_zero(i8 %x, i8 %y) {
+; CHECK-LABEL: @sdiv_abs_dividend_ge_abs_divisor_never_zero(
+; CHECK-NEXT:    [[XX:%.*]] = or i8 [[X:%.*]], -64
+; CHECK-NEXT:    [[YY:%.*]] = and i8 [[Y:%.*]], 63
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv i8 [[XX]], [[YY]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[DIV]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %xx = or i8 %x, 192  ; Ensure dividend is at least 192 or at most -64
+  %yy = and i8 %y, 63  ; Ensure divisor is at most 63 and at least -63
+  %div = sdiv i8 %xx, %yy
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+define i1 @udiv_exact_nonzero_dividend_never_zero(i8 %x) {
+; CHECK-LABEL: @udiv_exact_nonzero_dividend_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %num = or i8 %x, 1  ; Ensure dividend is non-zero
+  %denum = shl i8 1, 2 ; Divisor = 4
+  %div = udiv exact i8 %num, %denum
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}
+
+; Test exact division with non-zero dividend guarantees non-zero result for SDIV
+define i1 @sdiv_exact_nonzero_dividend_never_zero(i8 %x) {
+; CHECK-LABEL: @sdiv_exact_nonzero_dividend_never_zero(
+; CHECK-NEXT:    ret i1 true
+;
+  %num = or i8 %x, 1  ; Ensure dividend is non-zero
+  %denum = shl i8 1, 2 ; Divisor = 4
+  %div = sdiv exact i8 %num, %denum
+  %cmp = icmp ne i8 %div, 0
+  ret i1 %cmp
+}