[llvm] [VectorCombine] Don't fold non-idempotent shuffle reductions when shuffle duplicates element (PR #200778)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: lijinpei-amd

<details>
<summary>Changes</summary>

For odd vec-size, the shuffle scheme duplicates elements, making the transformation not equivalent when reduction-op non-idempotent.


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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+8) 
- (modified) llvm/test/Transforms/VectorCombine/fold-shuffle-chains-to-reduce.ll (+44-4) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index da8a52a6b3dca..b3ad21fe1115b 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -4002,6 +4002,7 @@ bool VectorCombine::foldShuffleChainsToReduce(Instruction &I) {
   InstWorklist.push(VecOpEE);
 
   bool IsPartialReduction = false;
+  bool HasLaneDuplication = false;
 
   while (!InstWorklist.empty()) {
     Value *CI = InstWorklist.front();
@@ -4121,6 +4122,7 @@ bool VectorCombine::foldShuffleChainsToReduce(Instruction &I) {
       // Update mask values.
       ShuffleMaskHalf *= 2;
       ShuffleMaskHalf -= (ExpectedParityMask & 1);
+      HasLaneDuplication |= (ExpectedParityMask & 1) != 0;
       ExpectedParityMask >>= 1;
 
       OrigCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
@@ -4152,6 +4154,12 @@ bool VectorCombine::foldShuffleChainsToReduce(Instruction &I) {
   if (ShouldBeCallOrBinInst && !IsPartialReduction)
     return false;
 
+  // If the parity masks duplicated any lane, the fold only preserves semantics
+  // for idempotent ops.
+  if (HasLaneDuplication && CommonBinOp &&
+      !Instruction::isIdempotent(*CommonBinOp))
+    return false;
+
   assert(VecSize != -1 && "Expected Match for Vector Size");
 
   Value *FinalVecV = PrevVecV[0];
diff --git a/llvm/test/Transforms/VectorCombine/fold-shuffle-chains-to-reduce.ll b/llvm/test/Transforms/VectorCombine/fold-shuffle-chains-to-reduce.ll
index e4cb41b238cf1..8111309dc5f53 100644
--- a/llvm/test/Transforms/VectorCombine/fold-shuffle-chains-to-reduce.ll
+++ b/llvm/test/Transforms/VectorCombine/fold-shuffle-chains-to-reduce.ll
@@ -47,11 +47,17 @@ define i16 @test_reduce_v3i16_and(<3 x i16> %a0) {
   ret i16 %5
 }
 
-define i16 @test_reduce_v6i16_xor(<6 x i16> %a0) {
-; CHECK-LABEL: define i16 @test_reduce_v6i16_xor(
+define i16 @test_no_reduce_v6i16_xor(<6 x i16> %a0) {
+; CHECK-LABEL: define i16 @test_no_reduce_v6i16_xor(
 ; CHECK-SAME: <6 x i16> [[A0:%.*]]) {
-; CHECK-NEXT:    [[TMP1:%.*]] = call i16 @llvm.vector.reduce.xor.v6i16(<6 x i16> [[A0]])
-; CHECK-NEXT:    ret i16 [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <6 x i16> [[A0]], <6 x i16> poison, <6 x i32> <i32 3, i32 4, i32 5, i32 poison, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP2:%.*]] = xor <6 x i16> [[A0]], [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <6 x i16> [[TMP2]], <6 x i16> poison, <6 x i32> <i32 1, i32 2, i32 poison, i32 poison, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP4:%.*]] = xor <6 x i16> [[TMP2]], [[TMP3]]
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <6 x i16> [[TMP4]], <6 x i16> poison, <6 x i32> <i32 1, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP6:%.*]] = xor <6 x i16> [[TMP4]], [[TMP5]]
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <6 x i16> [[TMP6]], i64 0
+; CHECK-NEXT:    ret i16 [[TMP7]]
 ;
   %1 = shufflevector <6 x i16> %a0, <6 x i16> poison, <6 x i32> <i32 3, i32 4, i32 5, i32 poison, i32 poison, i32 poison>
   %2 = xor <6 x i16> %a0, %1
@@ -63,6 +69,40 @@ define i16 @test_reduce_v6i16_xor(<6 x i16> %a0) {
   ret i16 %7
 }
 
+; Non-idempotent op (add) on a non-power-of-2 vector: can't fold into reduction.
+define i32 @test_no_reduce_v3i32_add(<3 x i32> %a0) {
+; CHECK-LABEL: define i32 @test_no_reduce_v3i32_add(
+; CHECK-SAME: <3 x i32> [[A0:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <3 x i32> [[A0]], <3 x i32> poison, <3 x i32> <i32 1, i32 2, i32 poison>
+; CHECK-NEXT:    [[TMP2:%.*]] = add <3 x i32> [[A0]], [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <3 x i32> [[TMP2]], <3 x i32> poison, <3 x i32> <i32 1, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP4:%.*]] = add <3 x i32> [[TMP2]], [[TMP3]]
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <3 x i32> [[TMP4]], i64 0
+; CHECK-NEXT:    ret i32 [[TMP5]]
+;
+  %1 = shufflevector <3 x i32> %a0, <3 x i32> poison, <3 x i32> <i32 1, i32 2, i32 poison>
+  %2 = add <3 x i32> %a0, %1
+  %3 = shufflevector <3 x i32> %2, <3 x i32> poison, <3 x i32> <i32 1, i32 poison, i32 poison>
+  %4 = add <3 x i32> %2, %3
+  %5 = extractelement <3 x i32> %4, i64 0
+  ret i32 %5
+}
+
+; Idempotent op (smax) on a non-power-of-2 vector: lane duplication is harmless
+define i16 @test_reduce_v3i16_smax(<3 x i16> %a0) {
+; CHECK-LABEL: define i16 @test_reduce_v3i16_smax(
+; CHECK-SAME: <3 x i16> [[A0:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i16 @llvm.vector.reduce.smax.v3i16(<3 x i16> [[A0]])
+; CHECK-NEXT:    ret i16 [[TMP1]]
+;
+  %1 = shufflevector <3 x i16> %a0, <3 x i16> poison, <3 x i32> <i32 1, i32 2, i32 poison>
+  %2 = tail call <3 x i16> @llvm.smax.v3i16(<3 x i16> %a0, <3 x i16> %1)
+  %3 = shufflevector <3 x i16> %2, <3 x i16> poison, <3 x i32> <i32 1, i32 poison, i32 poison>
+  %4 = tail call <3 x i16> @llvm.smax.v3i16(<3 x i16> %2, <3 x i16> %3)
+  %5 = extractelement <3 x i16> %4, i64 0
+  ret i16 %5
+}
+
 define i16 @test_reduce_v8i16_2(<8 x i16> %a0) {
 ; CHECK-LABEL: define i16 @test_reduce_v8i16_2(
 ; CHECK-SAME: <8 x i16> [[A0:%.*]]) {

``````````

</details>


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