From 9d6b92e29fee91c827aed67b699a8a4cc38dbaac Mon Sep 17 00:00:00 2001
From: Pranshu Goyal <78131793+pranshoe@users.noreply.github.com>
Date: Sat, 28 Mar 2026 22:20:19 +0530
Subject: [PATCH] [DAG] SelectionDAG::isKnownToBeAPowerOfTwo - add
 ISD::TRUNCATE handling and tests (#184365)

Closes #181654
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  5 +++
 .../CodeGen/AArch64/knownpow2-trunc-orzero.ll | 34 +++++++++++++++++++
 llvm/test/CodeGen/X86/known-pow2.ll           | 20 +++++------
 3 files changed, 48 insertions(+), 11 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/knownpow2-trunc-orzero.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 86afab7c4553..e53cef7b8be7 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4823,6 +4823,11 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
                                   Depth + 1);
   }
 
+  case ISD::TRUNCATE: {
+    return (OrZero || isKnownNeverZero(Val, Depth)) &&
+           isKnownToBeAPowerOfTwo(Val.getOperand(0), OrZero, Depth + 1);
+  }
+
   case ISD::ROTL:
   case ISD::ROTR:
     return isKnownToBeAPowerOfTwo(Val.getOperand(0), DemandedElts, OrZero,
diff --git a/llvm/test/CodeGen/AArch64/knownpow2-trunc-orzero.ll b/llvm/test/CodeGen/AArch64/knownpow2-trunc-orzero.ll
new file mode 100644
index 000000000000..b0c951537628
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/knownpow2-trunc-orzero.ll
@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc < %s -mtriple=arm64-unknown-unknown-gnu | FileCheck %s --check-prefix=CHECK
+
+define i8 @pow2_trunc_orzero_false_notzero(i32 %x, i8 %z){
+; CHECK-LABEL: pow2_trunc_orzero_false_notzero:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w1, #0xff
+; CHECK-NEXT:    and w9, w0, #0x7
+; CHECK-NEXT:    lsr w0, w8, w9
+; CHECK-NEXT:    ret
+
+  %amt = and i32 %x, 7
+  %y = shl i32 1, %amt
+  %tr = trunc i32 %y to i8
+  %div = udiv i8 %z, %tr
+  ret i8 %div
+}
+
+define i8 @pow2_trunc_orzero_false_maybezero(i32 %x, i8 %z){
+; CHECK-LABEL: pow2_trunc_orzero_false_maybezero:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    neg w8, w0
+; CHECK-NEXT:    and w9, w1, #0xff
+; CHECK-NEXT:    and w8, w0, w8
+; CHECK-NEXT:    and w8, w8, #0xff
+; CHECK-NEXT:    udiv w0, w9, w8
+; CHECK-NEXT:    ret
+
+  %neg_x = sub i32 0, %x
+  %y = and i32 %x, %neg_x
+  %tr = trunc i32 %y to i8
+  %div = udiv i8 %z, %tr
+  ret i8 %div
+}
diff --git a/llvm/test/CodeGen/X86/known-pow2.ll b/llvm/test/CodeGen/X86/known-pow2.ll
index 2806194cff76..5c395ec42350 100644
--- a/llvm/test/CodeGen/X86/known-pow2.ll
+++ b/llvm/test/CodeGen/X86/known-pow2.ll
@@ -1344,12 +1344,11 @@ define i32 @pow2_blsi_sub(i32 %x, i32 %a) {
 define i8 @pow2_trunc(i32 %x, i32 %a){
 ; CHECK-LABEL: pow2_trunc:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl %esi, %ecx
-; CHECK-NEXT:    negl %ecx
-; CHECK-NEXT:    andl %esi, %ecx
-; CHECK-NEXT:    movzbl %dil, %eax
-; CHECK-NEXT:    divb %cl
-; CHECK-NEXT:    movzbl %ah, %eax
+; CHECK-NEXT:    movl %esi, %eax
+; CHECK-NEXT:    negl %eax
+; CHECK-NEXT:    andl %esi, %eax
+; CHECK-NEXT:    decb %al
+; CHECK-NEXT:    andb %dil, %al
 ; CHECK-NEXT:    # kill: def $al killed $al killed $eax
 ; CHECK-NEXT:    retq
   %neg_a = sub i32 0, %a
@@ -1382,11 +1381,10 @@ define i8 @pow2_trunc_vec(<4 x i32> %x, <4 x i32> %a) {
 ; CHECK-NEXT:    pxor %xmm2, %xmm2
 ; CHECK-NEXT:    psubd %xmm1, %xmm2
 ; CHECK-NEXT:    pand %xmm1, %xmm2
-; CHECK-NEXT:    movd %xmm2, %ecx
-; CHECK-NEXT:    movd %xmm0, %eax
-; CHECK-NEXT:    movzbl %al, %eax
-; CHECK-NEXT:    divb %cl
-; CHECK-NEXT:    movzbl %ah, %eax
+; CHECK-NEXT:    pcmpeqd %xmm1, %xmm1
+; CHECK-NEXT:    paddb %xmm2, %xmm1
+; CHECK-NEXT:    pand %xmm0, %xmm1
+; CHECK-NEXT:    movd %xmm1, %eax
 ; CHECK-NEXT:    # kill: def $al killed $al killed $eax
 ; CHECK-NEXT:    retq
   %a.splat = shufflevector <4 x i32> %a, <4 x i32> poison, <4 x i32> zeroinitializer