[llvm] [Support] Adding fast-path for KnownBits's udiv when RHS is contant (PR #189779)

[Max Graey via llvm-commits]

https://github.com/MaxGraey updated https://github.com/llvm/llvm-project/pull/189779

>From bc71364b702b686d437e63bd8fb043ff3d001a6b Mon Sep 17 00:00:00 2001
From: MaxGraey <maxgraey at gmail.com>
Date: Wed, 1 Apr 2026 03:48:37 +0300
Subject: [PATCH 1/4] add fast-path for KnownBits::udiv when RHS is contant

---
 llvm/lib/Support/KnownBits.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
index 07e7781d0839d..5809fff3ad6ab 100644
--- a/llvm/lib/Support/KnownBits.cpp
+++ b/llvm/lib/Support/KnownBits.cpp
@@ -1209,6 +1209,20 @@ KnownBits KnownBits::udiv(const KnownBits &LHS, const KnownBits &RHS,
     return Known;
   }
 
+  if (RHS.isConstant()) {
+    const APInt Divisor = RHS.getConstant();
+    if (Divisor.isOne())
+      return LHS;
+
+    if (Divisor.isPowerOf2()) {
+      unsigned Shift = Divisor.logBase2();
+      KnownBits Known = LHS;
+      Known.One.lshrInPlace(Shift);
+      Known.Zero.lshrInPlace(Shift);
+      return Known;
+    }
+  }
+
   // We can figure out the minimum number of upper zero bits by doing
   // MaxNumerator / MinDenominator. If the Numerator gets smaller or Denominator
   // gets larger, the number of upper zero bits increases.

>From 98fbc0b1f97150ca44a9fe1b4ca9152665c75868 Mon Sep 17 00:00:00 2001
From: MaxGraey <maxgraey at gmail.com>
Date: Wed, 1 Apr 2026 03:56:48 +0300
Subject: [PATCH 2/4] refactor

---
 llvm/lib/Support/KnownBits.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
index 5809fff3ad6ab..25530d7e46727 100644
--- a/llvm/lib/Support/KnownBits.cpp
+++ b/llvm/lib/Support/KnownBits.cpp
@@ -1216,7 +1216,7 @@ KnownBits KnownBits::udiv(const KnownBits &LHS, const KnownBits &RHS,
 
     if (Divisor.isPowerOf2()) {
       unsigned Shift = Divisor.logBase2();
-      KnownBits Known = LHS;
+      Known = LHS;
       Known.One.lshrInPlace(Shift);
       Known.Zero.lshrInPlace(Shift);
       return Known;

>From f326a707d239257f17684715800cb511c523dda9 Mon Sep 17 00:00:00 2001
From: MaxGraey <maxgraey at gmail.com>
Date: Wed, 1 Apr 2026 05:01:20 +0300
Subject: [PATCH 3/4] fix: preserve high bits of LHS

---
 llvm/lib/Support/KnownBits.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/llvm/lib/Support/KnownBits.cpp b/llvm/lib/Support/KnownBits.cpp
index 25530d7e46727..56ec9de6af77e 100644
--- a/llvm/lib/Support/KnownBits.cpp
+++ b/llvm/lib/Support/KnownBits.cpp
@@ -1216,9 +1216,15 @@ KnownBits KnownBits::udiv(const KnownBits &LHS, const KnownBits &RHS,
 
     if (Divisor.isPowerOf2()) {
       unsigned Shift = Divisor.logBase2();
+
+      APInt MaxRes = LHS.getMaxValue().lshr(Shift);
+      unsigned LeadZ = MaxRes.countLeadingZeros();
+
       Known = LHS;
       Known.One.lshrInPlace(Shift);
       Known.Zero.lshrInPlace(Shift);
+      Known.Zero.setHighBits(LeadZ);
+
       return Known;
     }
   }

>From 7c6b11aeec5d951c899adbae80635abd721540e9 Mon Sep 17 00:00:00 2001
From: MaxGraey <maxgraey at gmail.com>
Date: Wed, 1 Apr 2026 11:31:59 +0300
Subject: [PATCH 4/4] update mir tests

---
 llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sdiv.mir | 4 ++--
 llvm/test/CodeGen/AArch64/GlobalISel/knownbits-udiv.mir | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sdiv.mir b/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sdiv.mir
index 7b315205bc2e8..ae56e9097c99f 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sdiv.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-sdiv.mir
@@ -8,7 +8,7 @@ body: |
   ; CHECK-LABEL: name: @Cst
   ; CHECK-NEXT: %0:_ KnownBits:01100100 SignBits:1
   ; CHECK-NEXT: %1:_ KnownBits:00000100 SignBits:5
-  ; CHECK-NEXT: %2:_ KnownBits:000????? SignBits:3
+  ; CHECK-NEXT: %2:_ KnownBits:00011001 SignBits:3
     %0:_(s8) = G_CONSTANT i8 100
     %1:_(s8) = G_CONSTANT i8 4
     %2:_(s8) = G_SDIV %0, %1
@@ -36,7 +36,7 @@ body: |
   ; CHECK-LABEL: name: @Exact
   ; CHECK-NEXT: %0:_ KnownBits:00001100 SignBits:4
   ; CHECK-NEXT: %1:_ KnownBits:00000100 SignBits:5
-  ; CHECK-NEXT: %2:_ KnownBits:000000?1 SignBits:6
+  ; CHECK-NEXT: %2:_ KnownBits:00000011 SignBits:6
     %0:_(s8) = G_CONSTANT i8 12
     %1:_(s8) = G_CONSTANT i8 4
     %2:_(s8) = exact G_SDIV %0, %1
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-udiv.mir b/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-udiv.mir
index dfbcd3f3c6188..0a0a00ce456ae 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-udiv.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/knownbits-udiv.mir
@@ -8,7 +8,7 @@ body: |
   ; CHECK-LABEL: name: @Cst
   ; CHECK-NEXT: %0:_ KnownBits:01100100 SignBits:1
   ; CHECK-NEXT: %1:_ KnownBits:00000100 SignBits:5
-  ; CHECK-NEXT: %2:_ KnownBits:000????? SignBits:3
+  ; CHECK-NEXT: %2:_ KnownBits:00011001 SignBits:3
     %0:_(s8) = G_CONSTANT i8 100
     %1:_(s8) = G_CONSTANT i8 4
     %2:_(s8) = G_UDIV %0, %1
@@ -36,7 +36,7 @@ body: |
   ; CHECK-LABEL: name: @Exact
   ; CHECK-NEXT: %0:_ KnownBits:00001100 SignBits:4
   ; CHECK-NEXT: %1:_ KnownBits:00000100 SignBits:5
-  ; CHECK-NEXT: %2:_ KnownBits:000000?1 SignBits:6
+  ; CHECK-NEXT: %2:_ KnownBits:00000011 SignBits:6
     %0:_(s8) = G_CONSTANT i8 12
     %1:_(s8) = G_CONSTANT i8 4
     %2:_(s8) = exact G_UDIV %0, %1