Bug type: crash
Component: LoopVectorize
Bug location: llvm/lib/Transforms/Vectorize/LoopVectorize.cpp: VPRecipeBuilder::collectScaledReductions
VPRecipeBuilder::getScaledReductions
Issue title: [LV] Assert "all accumulators in chain must have same scale factor"
Labels: vectorizers
crash-on-valid
Reported by: sjoerdmeijer

This IR:

```
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128-Fn32"
target triple = "aarch64-unknown-linux-gnu"

define void @_Z1fv() #0 {
entry:
  br label %for.body

for.cond.loopexit.loopexit:                       ; preds = %for.body
  ret void

for.body:                                         ; preds = %for.body, %entry
  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
  %add15.lcssa1921 = phi i32 [ 0, %entry ], [ %add.1, %for.body ]
  %conv8 = sext i16 0 to i32
  %add = add i32 %add15.lcssa1921, %conv8
  %conv8.1 = sext i16 0 to i32
  %add.1 = add i32 %add, %conv8.1
  store i32 %add.1, ptr null, align 4
  %indvars.iv.next = add i64 %indvars.iv, 1
  %0 = and i64 %indvars.iv.next, 255
  %exitcond.not = icmp eq i64 %0, 0
  br i1 %exitcond.not, label %for.cond.loopexit.loopexit, label %for.body
}

attributes #0 = { "target-cpu"="grace" }
```

Triggers this assert:

```
opt: /root/llvm-project/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp:8181: llvm::VPRecipeBase* llvm::VPRecipeBuilder::tryToCreatePartialReduction(llvm::VPInstruction*, unsigned int): Assertion `ScaleFactor == vputils::getVFScaleFactor(Accumulator->getDefiningRecipe()) && "all accumulators in chain must have same scale factor"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace and instructions to reproduce the bug.
Stack dump:
0.	Program arguments: /opt/compiler-explorer/clang-assertions-trunk/bin/opt -o /app/output.s -S -S -passes=loop-vectorize <source>
1.	Running pass "function(loop-vectorize<no-interleave-forced-only;no-vectorize-forced-only;>)" on module "<source>"
2.	Running pass "loop-vectorize<no-interleave-forced-only;no-vectorize-forced-only;>" on function "_Z1fv"
 #0 0x00000000059ad0e8 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/opt/compiler-explorer/clang-assertions-trunk/bin/opt+0x59ad0e8)
 #1 0x00000000059a9f94 SignalHandler(int, siginfo_t*, void*) Signals.cpp:0:0
 #2 0x00007efe89442520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
 #3 0x00007efe894969fc pthread_kill (/lib/x86_64-linux-gnu/libc.so.6+0x969fc)
 #4 0x00007efe89442476 gsignal (/lib/x86_64-linux-gnu/libc.so.6+0x42476)
 #5 0x00007efe894287f3 abort (/lib/x86_64-linux-gnu/libc.so.6+0x287f3)
 #6 0x00007efe8942871b (/lib/x86_64-linux-gnu/libc.so.6+0x2871b)
 #7 0x00007efe89439e96 (/lib/x86_64-linux-gnu/libc.so.6+0x39e96)
 #8 0x00000000036b2f1b llvm::VPRecipeBuilder::tryToCreatePartialReduction(llvm::VPInstruction*, unsigned int) (/opt/compiler-explorer/clang-assertions-trunk/bin/opt+0x36b2f1b)
 #9 0x00000000036b599f llvm::VPRecipeBuilder::tryToCreateWidenRecipe(llvm::VPSingleDefRecipe*, llvm::VFRange&) (/opt/compiler-explorer/clang-assertions-trunk/bin/opt+0x36b599f)
#10 0x00000000036d4e5e llvm::LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(std::unique_ptr<llvm::VPlan, std::default_delete<llvm::VPlan>>, llvm::VFRange&, llvm::LoopVersioning*) (/opt/compiler-explorer/clang-assertions-trunk/bin/opt+0x36d4e5e)
#11 0x00000000036d63a8 llvm::LoopVectorizationPlanner::buildVPlansWithVPRecipes(llvm::ElementCount, llvm::ElementCount) (/opt/compiler-explorer/clang-assertions-trunk/bin/opt+0x36d63a8)
```

Compile with: `opt -S -passes=loop-vectorize`, see also: https://godbolt.org/z/7oedf34Pq

Reproduced from this C reproducer:

```
int a, b;
extern char e;
int *c;
short *d;
void f() {
  for (int g;;)
    for (char h; h < 20; h++) {
      a = c[h];
      for (int i = 0; i < 22; i += 3) {
        e = 0;
        b += d[g];
      }
    }
}
```

Comments: author: sjoerdmeijer
body: CC: @sdesmalen-arm , @SamTebbs33

Was wondering if this is more fall out from the partial reductions work? 
author: sjoerdmeijer
body: This assert was introduced in 3ad5765e2341b by @fhahn 

Fix patch:
commit 67e35bbebbed5e33a173751d019c22e4effa65fd
Author: Florian Hahn <flo@fhahn.com>
Date:   Thu Nov 20 15:54:57 2025 +0000

    [LV] Check full partial reduction chains in order. (#168036)
    
    https://github.com/llvm/llvm-project/pull/162822 added another
    validation step to check if entries in a partial reduction chain have
    the same scale factor. But the validation was still dependent on the
    order of entries in PartialReductionChains, and would fail to reject
    some cases (e.g. if the first first link matched the scale of the second
    link, but the second link is invalidated later).
    
    To fix that, group chains by their starting phi nodes, then perform the
    validation for each chain, and if it fails, invalidate the whole chain
    for the phi.
    
    Fixes https://github.com/llvm/llvm-project/issues/167243.
    Fixes https://github.com/llvm/llvm-project/issues/167867.
    
    PR: https://github.com/llvm/llvm-project/pull/168036

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 8013c47f53db..741caab1fe58 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7992,13 +7992,15 @@ VPReplicateRecipe *VPRecipeBuilder::handleReplication(VPInstruction *VPI,
 /// Find all possible partial reductions in the loop and track all of those that
 /// are valid so recipes can be formed later.
 void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
-  // Find all possible partial reductions.
-  SmallVector<std::pair<PartialReductionChain, unsigned>>
-      PartialReductionChains;
-  for (const auto &[Phi, RdxDesc] : Legal->getReductionVars()) {
+  // Find all possible partial reductions, grouping chains by their PHI. This
+  // grouping allows invalidating the whole chain, if any link is not a valid
+  // partial reduction.
+  MapVector<Instruction *,
+            SmallVector<std::pair<PartialReductionChain, unsigned>>>
+      ChainsByPhi;
+  for (const auto &[Phi, RdxDesc] : Legal->getReductionVars())
     getScaledReductions(Phi, RdxDesc.getLoopExitInstr(), Range,
-                        PartialReductionChains);
-  }
+                        ChainsByPhi[Phi]);
 
   // A partial reduction is invalid if any of its extends are used by
   // something that isn't another partial reduction. This is because the
@@ -8006,8 +8008,9 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
 
   // Build up a set of partial reduction ops for efficient use checking.
   SmallPtrSet<User *, 4> PartialReductionOps;
-  for (const auto &[PartialRdx, _] : PartialReductionChains)
-    PartialReductionOps.insert(PartialRdx.ExtendUser);
+  for (const auto &[_, Chains] : ChainsByPhi)
+    for (const auto &[PartialRdx, _] : Chains)
+      PartialReductionOps.insert(PartialRdx.ExtendUser);
 
   auto ExtendIsOnlyUsedByPartialReductions =
       [&PartialReductionOps](Instruction *Extend) {
@@ -8018,31 +8021,38 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
 
   // Check if each use of a chain's two extends is a partial reduction
   // and only add those that don't have non-partial reduction users.
-  for (auto Pair : PartialReductionChains) {
-    PartialReductionChain Chain = Pair.first;
-    if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
-        (!Chain.ExtendB || ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
-      ScaledReductionMap.try_emplace(Chain.Reduction, Pair.second);
+  for (const auto &[_, Chains] : ChainsByPhi) {
+    for (const auto &[Chain, Scale] : Chains) {
+      if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
+          (!Chain.ExtendB ||
+           ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
+        ScaledReductionMap.try_emplace(Chain.Reduction, Scale);
+    }
   }
 
   // Check that all partial reductions in a chain are only used by other
   // partial reductions with the same scale factor. Otherwise we end up creating
   // users of scaled reductions where the types of the other operands don't
   // match.
-  for (const auto &[Chain, Scale] : PartialReductionChains) {
-    auto AllUsersPartialRdx = [ScaleVal = Scale, this](const User *U) {
-      auto *UI = cast<Instruction>(U);
-      if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
-        return all_of(UI->users(), [ScaleVal, this](const User *U) {
-          auto *UI = cast<Instruction>(U);
-          return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal;
-        });
+  for (const auto &[Phi, Chains] : ChainsByPhi) {
+    for (const auto &[Chain, Scale] : Chains) {
+      auto AllUsersPartialRdx = [ScaleVal = Scale, RdxPhi = Phi,
+                                 this](const User *U) {
+        auto *UI = cast<Instruction>(U);
+        if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader())
+          return UI == RdxPhi;
+        return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
+               !OrigLoop->contains(UI->getParent());
+      };
+
+      // If any partial reduction entry for the phi is invalid, invalidate the
+      // whole chain.
+      if (!all_of(Chain.Reduction->users(), AllUsersPartialRdx)) {
+        for (const auto &[Chain, _] : Chains)
+          ScaledReductionMap.erase(Chain.Reduction);
+        break;
       }
-      return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
-             !OrigLoop->contains(UI->getParent());
-    };
-    if (!all_of(Chain.Reduction->users(), AllUsersPartialRdx))
-      ScaledReductionMap.erase(Chain.Reduction);
+    }
   }
 }
 


Regression tests:
file: llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll
commands: opt --mattr=+neon,+dotprod -passes=loop-vectorize -force-vector-interleave=1 -enable-epilogue-vectorization=false -S %s
tests: test_name: chained_sext_adds
test_body: target triple = "arm64-apple-macosx"

define void @chained_sext_adds(ptr noalias %src, ptr noalias %dst) #0 {
entry:
  br label %loop

exit:                                             ; preds = %loop
  ret void

loop:                                             ; preds = %loop, %entry
  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
  %red = phi i32 [ 0, %entry ], [ %add.1, %loop ]
  %gep.src = getelementptr i8, ptr %src, i64 %iv
  %l = load i8, ptr %gep.src, align 1
  %conv8 = sext i8 %l to i32
  %add = add i32 %red, %conv8
  %conv8.1 = sext i8 %l to i32
  %add.1 = add i32 %add, %conv8.1
  %gep.dst = getelementptr i8, ptr %dst, i64 %iv
  store i32 %add.1, ptr %dst, align 4
  %iv.next = add i64 %iv, 1
  %exitcond = icmp eq i64 %iv.next, 1000
  br i1 %exitcond, label %exit, label %loop
}

attributes #0 = { "target-cpu"="grace" }