[LV] Check all users of partial reductions in chain have same scale.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7915,6 +7915,26 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
  (!Chain.ExtendB || ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
  ScaledReductionMap.try_emplace(Chain.Reduction, Pair.second);
  }
+
+ // 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;
+ });
+ }
+ return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
+ !OrigLoop->contains(UI->getParent());
+ };
+ if (!all_of(Chain.Reduction->users(), AllUsersPartialRdx))
+ ScaledReductionMap.erase(Chain.Reduction);
+ }
 }
 
 bool VPRecipeBuilder::getScaledReductions(
@@ -8098,11 +8118,8 @@ VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(VPSingleDefRecipe *R,
  if (isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
  return tryToWidenMemory(Instr, Operands, Range);
 
- if (std::optional<unsigned> ScaleFactor = getScalingForReduction(Instr)) {
- if (auto PartialRed =
- tryToCreatePartialReduction(Instr, Operands, ScaleFactor.value()))
- return PartialRed;
- }
+ if (std::optional<unsigned> ScaleFactor = getScalingForReduction(Instr))
+ return tryToCreatePartialReduction(Instr, Operands, ScaleFactor.value());
 
  if (!shouldWiden(Instr, Range))
  return nullptr;
@@ -8136,9 +8153,9 @@ VPRecipeBuilder::tryToCreatePartialReduction(Instruction *Reduction,
  isa<VPPartialReductionRecipe>(BinOpRecipe))
  std::swap(BinOp, Accumulator);
 
- if (ScaleFactor !=
- vputils::getVFScaleFactor(Accumulator->getDefiningRecipe()))
- return nullptr;
+ assert(ScaleFactor ==
+  vputils::getVFScaleFactor(Accumulator->getDefiningRecipe()) &&
+  "all accumulators in chain must have same scale factor");
 
  unsigned ReductionOpcode = Reduction->getOpcode();
  if (ReductionOpcode == Instruction::Sub) {

llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll

@@ -70,3 +70,28 @@ loop:
 exit:
  ret i32 %red.next
 }
+
+define i16 @test_incomplete_chain_without_mul(ptr noalias %dst, ptr %A, ptr %B) #0 {
+entry:
+ br label %loop
+
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %red = phi i16 [ 0, %entry ], [ %red.next, %loop ]
+ %l.a = load i8, ptr %A, align 1
+ %a.ext = zext i8 %l.a to i16
+ store i16 %a.ext, ptr %dst, align 2
+ %l.b = load i8, ptr %B, align 1
+ %b.ext = zext i8 %l.b to i16
+ %add = add i16 %red, %b.ext
+ %add.1 = add i16 %add, %a.ext
+ %red.next = add i16 %add.1, %b.ext
+ %iv.next = add i64 %iv, 1
+ %ec = icmp ult i64 %iv, 1024
+ br i1 %ec, label %loop, label %exit
+
+exit:
+ ret i16 %red.next
+}
+
+attributes #0 = { "target-cpu"="grace" }