[llvm] 3db0f8c - [SLP]Update TrackedToOrig mappings after reduction vectorization

[Alexey Bataev via llvm-commits]

Author: Alexey Bataev
Date: 2024-09-23T10:52:03-07:00
New Revision: 3db0f8c895d4e814a18b754f9afbb1e03bd839a5

URL: https://github.com/llvm/llvm-project/commit/3db0f8c895d4e814a18b754f9afbb1e03bd839a5
DIFF: https://github.com/llvm/llvm-project/commit/3db0f8c895d4e814a18b754f9afbb1e03bd839a5.diff

LOG: [SLP]Update TrackedToOrig mappings after reduction vectorization

Need to update mappings in TrackedToOrig to correctly provide mapping
between updated reduced value after vectorization and its original
value, otherwise the compiler might miss this update and it may cause
compiler crash later, when it tries to find the original instruction
mapping for the updated value.

Fixes https://github.com/llvm/llvm-project/issues/109376

Added: 
    llvm/test/Transforms/SLPVectorizer/X86/reduced-value-vectorized-later.ll

Modified: 
    llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index a88702b81096ed..04b8fc09a724dd 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -18668,6 +18668,14 @@ class HorizontalReduction {
         // Vectorize a tree.
         Value *VectorizedRoot =
             V.vectorizeTree(LocalExternallyUsedValues, InsertPt);
+        // Update TrackedToOrig mapping, since the tracked values might be
+        // updated.
+        for (Value *RdxVal : Candidates) {
+          Value *OrigVal = TrackedToOrig.at(RdxVal);
+          Value *TransformedRdxVal = TrackedVals.at(OrigVal);
+          if (TransformedRdxVal != RdxVal)
+            TrackedToOrig.try_emplace(TransformedRdxVal, OrigVal);
+        }
 
         Builder.SetInsertPoint(InsertPt);
 

diff  --git a/llvm/test/Transforms/SLPVectorizer/X86/reduced-value-vectorized-later.ll b/llvm/test/Transforms/SLPVectorizer/X86/reduced-value-vectorized-later.ll
new file mode 100644
index 00000000000000..6b0b22b90510c1
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/reduced-value-vectorized-later.ll
@@ -0,0 +1,41 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=slp-vectorizer -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+
+define i16 @test() {
+; CHECK-LABEL: define i16 @test() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call i16 @llvm.vector.reduce.or.v8i16(<8 x i16> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call i16 @llvm.vector.reduce.or.v4i16(<4 x i16> zeroinitializer)
+; CHECK-NEXT:    [[OP_RDX:%.*]] = or i16 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    [[OP_RDX1:%.*]] = or i16 [[OP_RDX]], 0
+; CHECK-NEXT:    ret i16 [[OP_RDX1]]
+;
+entry:
+  %subi = add i16 0, 0
+  %sub40.i = add i16 %subi, 0
+  %sub41.i = add i16 %subi, 0
+  %sub42.i = add i16 %subi, 0
+  %sub43.i = add i16 %subi, 0
+  %sub44.i = add i16 %subi, 0
+  %sub45.i = add i16 %subi, 0
+  %sub46.i = add i16 0, 0
+  %sub47.i = add i16 0, 0
+  %sub48.i = add i16 0, 0
+  %sub49.i = add i16 0, 0
+  %or40.i = or i16 %sub40.i, %sub41.i
+  %or41.i = or i16 %or40.i, %sub42.i
+  %or42.i = or i16 %or41.i, %sub43.i
+  %or43.i = or i16 %or42.i, %sub44.i
+  %or44.i = or i16 %or43.i, %sub45.i
+  %or45.i = or i16 %or44.i, %sub46.i
+  %or46.i = or i16 %or45.i, %sub47.i
+  %or47.i = or i16 %or46.i, %sub48.i
+  %or48.i = or i16 %or47.i, %sub49.i
+  %or50.i = or i16 %or48.i, %subi
+  %subii = add i16 0, 0
+  %subi16.i = add i16 %subii, 0
+  %subi17.i = add i16 %subii, 0
+  %0 = or i16 %subi16.i, %subi17.i
+  %1 = or i16 %0, %or50.i
+  ret i16 %1
+}