[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 d1c82b241a1ed5de9070f641b949a25e97dbe937 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-powof2-shuffle.ll | 12 ++++++++++
 2 files changed, 36 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-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
+}