[llvm] [DAG] isKnownToBeAPowerOfTwo - add ISD::VECTOR_SHUFFLE handling (PR #185203)

[Arjun Parmar via llvm-commits]

https://github.com/akparmar004 updated https://github.com/llvm/llvm-project/pull/185203

>From 18161b6a584800516423a9df25cd6d9b1f53c411 Mon Sep 17 00:00:00 2001
From: akparmar004 <akparmar0404 at gmail.com>
Date: Sat, 7 Mar 2026 22:21:33 +0530
Subject: [PATCH] [DAG] isKnownToBeAPowerOfTwo - add ISD::VECTOR_SHUFFLE
 handling

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 24 +++++++++++++
 llvm/test/CodeGen/X86/known-pow2.ll           | 34 +++++++++++++++++++
 llvm/test/CodeGen/X86/known-powof2-shuffle.ll | 12 +++++++
 3 files changed, 70 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/known-powof2-shuffle.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index a2318fc034fa2..4ec771d7fd41f 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4784,6 +4784,30 @@ bool SelectionDAG::isKnownToBeAPowerOfTwo(SDValue Val,
     // vscale(power-of-two) is a power-of-two
     return isKnownToBeAPowerOfTwo(Val.getOperand(0), /*OrZero=*/false,
                                   Depth + 1);
+
+  case ISD::VECTOR_SHUFFLE: {
+    assert(!Val.getValueType().isScalableVector());
+    // Demanded elements with undef shuffle mask elements are unknown
+    // - we cannot guarantee they are a power of two, so return false.
+    APInt DemandedLHS, DemandedRHS;
+    const ShuffleVectorSDNode *SVN = cast<ShuffleVectorSDNode>(Val);
+    assert(NumElts == SVN->getMask().size() && "Unexpected vector size");
+    if (!getShuffleDemandedElts(NumElts, SVN->getMask(), DemandedElts,
+                                DemandedLHS, DemandedRHS))
+      return false;
+
+    // All demanded elements from LHS must be known power of two.
+    if (!!DemandedLHS && !isKnownToBeAPowerOfTwo(Val.getOperand(0), DemandedLHS,
+                                                 OrZero, Depth + 1))
+      return false;
+
+    // All demanded elements from RHS must be known power of two.
+    if (!!DemandedRHS && !isKnownToBeAPowerOfTwo(Val.getOperand(1), DemandedRHS,
+                                                 OrZero, Depth + 1))
+      return false;
+
+    return true;
+  }
   }
 
   // More could be done here, though the above checks are enough
diff --git a/llvm/test/CodeGen/X86/known-pow2.ll b/llvm/test/CodeGen/X86/known-pow2.ll
index a9b28458e9518..2fccbb088fb98 100644
--- a/llvm/test/CodeGen/X86/known-pow2.ll
+++ b/llvm/test/CodeGen/X86/known-pow2.ll
@@ -1241,3 +1241,37 @@ define i32 @pow2_rotr_extract_vec(<4 x i32> %a0, <4 x i32> %rotamt, i32 %x, ptr
   %res = urem i32 %x, %elt
   ret i32 %res
 }
+
+define <4 x i32> @pow2_shuffle_vec(<4 x i32> %a0, <4 x i32> %a1, <4 x i32> %a2) {
+; CHECK-LABEL: pow2_shuffle_vec:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pxor %xmm3, %xmm3
+; CHECK-NEXT:    pxor %xmm4, %xmm4
+; CHECK-NEXT:    pcmpgtd %xmm0, %xmm4
+; CHECK-NEXT:    pcmpgtd %xmm1, %xmm3
+; CHECK-NEXT:    movdqa %xmm4, %xmm0
+; CHECK-NEXT:    pandn {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0
+; CHECK-NEXT:    pand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm4
+; CHECK-NEXT:    por %xmm0, %xmm4
+; CHECK-NEXT:    movdqa %xmm3, %xmm0
+; CHECK-NEXT:    pandn {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0
+; CHECK-NEXT:    pand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm3
+; CHECK-NEXT:    por %xmm0, %xmm3
+; CHECK-NEXT:    pshufd {{.*#+}} xmm0 = xmm3[1,3,2,3]
+; CHECK-NEXT:    pshufd {{.*#+}} xmm1 = xmm4[0,2,2,3]
+; CHECK-NEXT:    punpckldq {{.*#+}} xmm1 = xmm1[0],xmm0[0],xmm1[1],xmm0[1]
+; CHECK-NEXT:    pcmpeqd %xmm0, %xmm0
+; CHECK-NEXT:    paddd %xmm1, %xmm0
+; CHECK-NEXT:    pand %xmm2, %xmm0
+; CHECK-NEXT:    retq
+  %cmp0 = icmp sgt <4 x i32> zeroinitializer, %a0
+  %cmp1 = icmp sgt <4 x i32> zeroinitializer, %a1
+  %sel0 = select <4 x i1> %cmp0, <4 x i32> <i32 4, i32 4, i32 4, i32 4>,
+                                  <4 x i32> <i32 8, i32 8, i32 8, i32 8>
+  %sel1 = select <4 x i1> %cmp1, <4 x i32> <i32 2, i32 2, i32 2, i32 2>,
+                                  <4 x i32> <i32 16, i32 16, i32 16, i32 16>
+  %shuf = shufflevector <4 x i32> %sel0, <4 x i32> %sel1,
+                        <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+  %res = urem <4 x i32> %a2, %shuf
+  ret <4 x i32> %res
+}
diff --git a/llvm/test/CodeGen/X86/known-powof2-shuffle.ll b/llvm/test/CodeGen/X86/known-powof2-shuffle.ll
new file mode 100644
index 0000000000000..fff81fecb879e
--- /dev/null
+++ b/llvm/test/CodeGen/X86/known-powof2-shuffle.ll
@@ -0,0 +1,12 @@
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+
+; Test that isKnownToBeAPowerOfTwo handles VECTOR_SHUFFLE
+
+; A shuffle of two power-of-2 vectors should be recognized as power-of-2
+define <4 x i32> @shuffle_pow2(<4 x i32> %a) {
+; CHECK-LABEL: shuffle_pow2:
+  %pow2 = and <4 x i32> %a, <i32 4, i32 4, i32 4, i32 4>
+  %splat = shufflevector <4 x i32> %pow2, <4 x i32> poison,
+                         <4 x i32> <i32 0, i32 0, i32 0, i32 0>
+  ret <4 x i32> %splat
+}