[llvm] [SLP]Fix reused scales for sub-fields gathered in natural order (PR #226828)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Alexey Bataev (alexey-bataev)

<details>
<summary>Changes</summary>

The reduction root sub-fields were emitted in the natural order, but the
reused values were scaled by the original lane order. Reorder the root
scalars instead.

Fixes https://github.com/llvm/llvm-project/pull/224919#issuecomment-5858504241


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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp (+16-14) 
- (modified) llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll (+2-2) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 0f1fbea086abc..442cc9a0895b2 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -14249,6 +14249,22 @@ void BoUpSLP::transformNodes() {
     });
   };
 
+  // The reduction root gather feeds only the commutative reduction, which
+  // maps the reduced values to the lanes by the root scalars. Put the distinct
+  // sub-fields of the same wider scalar in the natural order to emit them
+  // without a permutation.
+  if (UserIgnoreList && getRootNode().isGather() &&
+      getRootNode().ReuseShuffleIndices.empty())
+    if (std::optional<std::tuple<Value *, unsigned, SmallVector<int>>> Fields =
+            matchGatheredExtractedFields(getRootNode().Scalars, *DL)) {
+      auto &[Src, FieldWidth, Mask] = *Fields;
+      SmallVector<int> SortedMask(Mask);
+      sort(SortedMask);
+      if (equal(SortedMask,
+                seq<int>(Src->getType()->getIntegerBitWidth() / FieldWidth)))
+        reorderScalars(getRootNode().Scalars, Mask);
+    }
+
   // Try to reorder gather nodes for better vectorization opportunities.
   for (unsigned Idx : seq<unsigned>(BaseGraphSize)) {
     TreeEntry &E = *VectorizableTree[Idx];
@@ -22792,20 +22808,6 @@ std::optional<ResTy> BoUpSLP::processExtractedFieldsGather(
   Value *Src = std::get<0>(*ExtractedFields);
   unsigned FieldWidth = std::get<1>(*ExtractedFields);
   SmallVector<int> &Mask = std::get<2>(*ExtractedFields);
-  unsigned NumFields = Src->getType()->getIntegerBitWidth() / FieldWidth;
-  // The vector of the reduction root gather feeds only the commutative
-  // reduction; other users of the gathered scalars keep using the scalars,
-  // so the lane order is unobservable. Emit the fields in the natural order
-  // and skip the permutation when each lane holds a distinct field.
-  if (!E->UserTreeIndex && UserIgnoreList && Mask.size() == NumFields) {
-    SmallVector<int> SortedMask(Mask);
-    sort(SortedMask);
-    // Each field is held at most once; poison lanes are ignored.
-    if (adjacent_find(SortedMask, [](int A, int B) {
-          return A != PoisonMaskElem && A == B;
-        }) == SortedMask.end())
-      std::iota(Mask.begin(), Mask.end(), 0);
-  }
   return ShuffleBuilder.createExtractedFieldsVector(Src, FieldWidth, Mask, *E);
 }
 
diff --git a/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll b/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll
index 68651173527f3..e682309e1aa9d 100644
--- a/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll
+++ b/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll
@@ -436,7 +436,7 @@ define i32 @sum4_i32_reused(i32 %x) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32 [[X]] to <4 x i8>
 ; CHECK-NEXT:    [[TMP2:%.*]] = zext <4 x i8> [[TMP1]] to <4 x i16>
 ; CHECK-NEXT:    [[TMP3:%.*]] = zext <4 x i16> [[TMP2]] to <4 x i32>
-; CHECK-NEXT:    [[TMP4:%.*]] = mul <4 x i32> [[TMP3]], <i32 3, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP4:%.*]] = mul <4 x i32> [[TMP3]], <i32 1, i32 1, i32 1, i32 3>
 ; CHECK-NEXT:    [[TMP5:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[TMP4]])
 ; CHECK-NEXT:    ret i32 [[TMP5]]
 ;
@@ -467,7 +467,7 @@ define i16 @xor8_i64_reused(i64 %x) {
 ; CHECK-SAME: i64 [[X:%.*]]) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i64 [[X]] to <8 x i8>
 ; CHECK-NEXT:    [[TMP2:%.*]] = zext <8 x i8> [[TMP1]] to <8 x i16>
-; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <8 x i16> [[TMP2]], <8 x i16> zeroinitializer, <8 x i32> <i32 8, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <8 x i16> [[TMP2]], <8 x i16> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 8>
 ; CHECK-NEXT:    [[TMP4:%.*]] = call i16 @llvm.vector.reduce.xor.v8i16(<8 x i16> [[TMP3]])
 ; CHECK-NEXT:    ret i16 [[TMP4]]
 ;

``````````

</details>


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