[llvm] [SLP]Extract vectorized source of gathered sub-fields (PR #226437)

[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 bitcast emitted for gathered sub-fields used the source scalar
directly, even if it was vectorized and erased, causes the crash,
register the bitcast as an external user of the vectorized source.

Fixes the crash reported in #<!-- -->224919.


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


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp (+8) 
- (modified) llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll (+38) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index b76cf0677c2f6..9179562e6dbf3 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -22583,6 +22583,14 @@ class BoUpSLP::ShuffleInstructionBuilder final : public BaseShuffleAnalysis {
     if (auto *I = dyn_cast<Instruction>(Vec)) {
       R.GatherShuffleExtractSeq.insert(I);
       R.CSEBlocks.insert(I->getParent());
+      // A vectorized source scalar is erased; extract it for the bitcast.
+      ArrayRef<TreeEntry *> TEs = R.getTreeEntries(Src);
+      const auto *It = find_if_not(TEs, [&](const TreeEntry *TE) {
+        return R.TransformedToGatherNodes.contains(TE) ||
+               R.DeletedNodes.contains(TE);
+      });
+      if (It != TEs.end())
+        R.ExternalUses.emplace_back(Src, I, **It, (*It)->findLaneForValue(Src));
     }
     Vec = createShuffle(Vec, /*V2=*/nullptr, Mask);
     // The extracted fields are unsigned.
diff --git a/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll b/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll
index e10b2f1899848..1c43215b77d1b 100644
--- a/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll
+++ b/llvm/test/Transforms/SLPVectorizer/X86/extracted-subfields.ll
@@ -388,3 +388,41 @@ start:
   store i16 %b3, ptr %out, align 2
   ret i16 %a7
 }
+
+define void @vectorized_src(ptr %p, ptr %a, i64 %x) {
+; CHECK-LABEL: define void @vectorized_src(
+; CHECK-SAME: ptr [[P:%.*]], ptr [[A:%.*]], i64 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds nuw i8, ptr [[P]], i64 8
+; CHECK-NEXT:    [[TMP0:%.*]] = load <2 x i64>, ptr [[P]], align 8
+; CHECK-NEXT:    [[LD1:%.*]] = load i64, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i64 [[LD1]] to <2 x i32>
+; CHECK-NEXT:    [[TMP2:%.*]] = sitofp <2 x i32> [[TMP1]] to <2 x double>
+; CHECK-NEXT:    [[TMP3:%.*]] = trunc <2 x i64> [[TMP0]] to <2 x i32>
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc i64 [[X]] to i32
+; CHECK-NEXT:    [[TMP5:%.*]] = insertelement <2 x i32> [[TMP3]], i32 [[TMP4]], i64 1
+; CHECK-NEXT:    [[TMP6:%.*]] = sitofp <2 x i32> [[TMP5]] to <2 x double>
+; CHECK-NEXT:    [[TMP7:%.*]] = fsub <2 x double> [[TMP2]], [[TMP6]]
+; CHECK-NEXT:    store <2 x double> [[TMP7]], ptr [[A]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  %ld0 = load i64, ptr %p, align 8
+  %gep = getelementptr inbounds nuw i8, ptr %p, i64 8
+  %ld1 = load i64, ptr %gep, align 8
+  %t0 = trunc i64 %ld1 to i32
+  %c0 = sitofp i32 %t0 to double
+  %t1 = trunc i64 %ld0 to i32
+  %c1 = sitofp i32 %t1 to double
+  %sub0 = fsub double %c0, %c1
+  store double %sub0, ptr %a, align 8
+  %s = lshr i64 %ld1, 32
+  %t2 = trunc nuw i64 %s to i32
+  %c2 = sitofp i32 %t2 to double
+  %t3 = trunc i64 %x to i32
+  %c3 = sitofp i32 %t3 to double
+  %sub1 = fsub double %c2, %c3
+  %gep.a = getelementptr inbounds nuw i8, ptr %a, i64 8
+  store double %sub1, ptr %gep.a, align 8
+  ret void
+}

``````````

</details>


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