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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-selectiondag

Author: Arjun Parmar (akparmar004)

<details>
<summary>Changes</summary>

closes #<!-- -->183350

---
Full diff: https://github.com/llvm/llvm-project/pull/185203.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+24) 
- (added) llvm/test/CodeGen/X86/known-powof2-shuffle.ll (+10) 


``````````diff
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..7b2dd24baf6d8
--- /dev/null
+++ b/llvm/test/CodeGen/X86/known-powof2-shuffle.ll
@@ -0,0 +1,10 @@
+; 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:
+  %splat = shufflevector <4 x i32> %a, <4 x i32> undef, <4 x i32> <i32 0, i32 0, i32 0, i32 0>
+  ret <4 x i32> %splat
+}

``````````

</details>


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