[VPlan] Handle live-in extend operands in partial reduction ::computeCost #163175

fhahn · 2025-10-13T10:52:52Z

In some cases, extend operands can be live-ins. Use dyn_cast_if_present, to handle this case without crashing.

Fixes #162902.

llvmbot · 2025-10-13T10:53:30Z

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-vectorizers

Author: Florian Hahn (fhahn)

Changes

In some cases, extend operands can be live-ins. Use dyn_cast_if_present, to handle this case without crashing.

Fixes #162902.

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

2 Files Affected:

(modified) llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp (+4-2)
(modified) llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-constant-ops.ll (+58)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 2368d18b0373c..95d220d53a2a9 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -330,6 +330,8 @@ VPPartialReductionRecipe::computeCost(ElementCount VF,
   auto HandleWiden = [&](VPWidenRecipe *Widen) {
     if (match(Widen, m_Sub(m_ZeroInt(), m_VPValue(Op)))) {
       Widen = dyn_cast<VPWidenRecipe>(Op->getDefiningRecipe());
+      if (!Widen)
+        return;
     }
     Opcode = Widen->getOpcode();
     VPRecipeBase *ExtAR = Widen->getOperand(0)->getDefiningRecipe();
@@ -355,10 +357,10 @@ VPPartialReductionRecipe::computeCost(ElementCount VF,
     ExtAType = GetExtendKind(OpR);
   } else if (isa<VPReductionPHIRecipe>(OpR)) {
     auto RedPhiOp1R = getOperand(1)->getDefiningRecipe();
-    if (isa<VPWidenCastRecipe>(RedPhiOp1R)) {
+    if (isa_and_nonnull<VPWidenCastRecipe>(RedPhiOp1R)) {
       InputTypeA = Ctx.Types.inferScalarType(RedPhiOp1R->getOperand(0));
       ExtAType = GetExtendKind(RedPhiOp1R);
-    } else if (auto Widen = dyn_cast<VPWidenRecipe>(RedPhiOp1R))
+    } else if (auto Widen = dyn_cast_if_present<VPWidenRecipe>(RedPhiOp1R))
       HandleWiden(Widen);
   } else if (auto Widen = dyn_cast<VPWidenRecipe>(OpR)) {
     HandleWiden(Widen);
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-constant-ops.ll b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-constant-ops.ll
index b033f6051f812..bf6967fecbc92 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-constant-ops.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-constant-ops.ll
@@ -467,3 +467,61 @@ loop:
 exit:
   ret i32 %red.next
 }
+
+; Test case for https://github.com/llvm/llvm-project/issues/162902.
+define void @partial_reduction_zext_const(i64 %arg, ptr %ptr) {
+; CHECK-LABEL: define void @partial_reduction_zext_const(
+; CHECK-SAME: i64 [[ARG:%.*]], ptr [[PTR:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = sub i64 100, [[ARG]]
+; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 [[TMP0]], 4
+; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], label %[[SCALAR_PH:.*]], label %[[VECTOR_PH:.*]]
+; CHECK:       [[VECTOR_PH]]:
+; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i64 [[TMP0]], 4
+; CHECK-NEXT:    [[N_VEC:%.*]] = sub i64 [[TMP0]], [[N_MOD_VF]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i64 [[ARG]], [[N_VEC]]
+; CHECK-NEXT:    br label %[[VECTOR_BODY:.*]]
+; CHECK:       [[VECTOR_BODY]]:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <4 x i8> [ zeroinitializer, %[[VECTOR_PH]] ], [ [[TMP2:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP2]] = add <4 x i8> [[VEC_PHI]], splat (i8 2)
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[TMP3]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP16:![0-9]+]]
+; CHECK:       [[MIDDLE_BLOCK]]:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i8 @llvm.vector.reduce.add.v4i8(<4 x i8> [[TMP2]])
+; CHECK-NEXT:    store i8 [[TMP4]], ptr [[PTR]], align 1
+; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 [[TMP0]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[CMP_N]], label %[[EXIT:.*]], label %[[SCALAR_PH]]
+; CHECK:       [[SCALAR_PH]]:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[TMP1]], %[[MIDDLE_BLOCK]] ], [ [[ARG]], %[[ENTRY]] ]
+; CHECK-NEXT:    [[BC_MERGE_RDX:%.*]] = phi i8 [ [[TMP4]], %[[MIDDLE_BLOCK]] ], [ 0, %[[ENTRY]] ]
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], %[[SCALAR_PH]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[PARTIAL:%.*]] = phi i8 [ [[BC_MERGE_RDX]], %[[SCALAR_PH]] ], [ [[PARTIAL_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[ZERO_EXT:%.*]] = zext i2 -2 to i8
+; CHECK-NEXT:    [[PARTIAL_NEXT]] = add i8 [[PARTIAL]], [[ZERO_EXT]]
+; CHECK-NEXT:    store i8 [[PARTIAL_NEXT]], ptr [[PTR]], align 1
+; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV_NEXT]], 100
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT]], label %[[LOOP]], !llvm.loop [[LOOP17:![0-9]+]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i64 [ %arg, %entry ], [ %iv.next, %loop ]
+  %partial = phi i8 [ 0, %entry ], [ %partial.next, %loop ]
+  %zero.ext = zext i2 2 to i8
+  %partial.next = add i8 %partial, %zero.ext
+  store i8 %partial.next, ptr %ptr, align 1
+  %iv.next = add i64 %iv, 1
+  %ec = icmp eq i64 %iv.next, 100
+  br i1 %ec, label %exit, label %loop
+
+exit:
+  ret void
+}

sjoerdmeijer

Thanks for fixing, looks like a good fix to me.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

david-arm · 2025-10-14T13:55:40Z

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

    if (match(Widen, m_Sub(m_ZeroInt(), m_VPValue(Op)))) {
      Widen = dyn_cast<VPWidenRecipe>(Op->getDefiningRecipe());
+      if (!Widen)
+        return;


What if Op is not a live-in - does it matter? If should never be any other type of recipe, then perhaps worth asserting it's a live-in?

Currently it must always be a live-in I think. I added an assert for now, thanks

SamTebbs33 · 2025-10-14T15:32:07Z

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

  } else if (isa<VPReductionPHIRecipe>(OpR)) {
    auto RedPhiOp1R = getOperand(1)->getDefiningRecipe();
-    if (isa<VPWidenCastRecipe>(RedPhiOp1R)) {
+    if (isa_and_nonnull<VPWidenCastRecipe>(RedPhiOp1R)) {
      InputTypeA = Ctx.Types.inferScalarType(RedPhiOp1R->getOperand(0));
      ExtAType = GetExtendKind(RedPhiOp1R);
-    } else if (auto Widen = dyn_cast<VPWidenRecipe>(RedPhiOp1R))
+    } else if (auto Widen = dyn_cast_if_present<VPWidenRecipe>(RedPhiOp1R))
      HandleWiden(Widen);
  } else if (auto Widen = dyn_cast<VPWidenRecipe>(OpR)) {


Would you mind making the same changes to the other isa checks and dyn_cast calls that use OpR in this function? I've noticed they are still problematic with constants while working on partial reduction bundling.

Sure, I can do that, assuming you mean the checks below? Do you think there are any test cases that could hit the code, without your changes?

…Cost In some cases, extend operands can be live-ins. Use dyn_cast_if_present, to handle this case without crashing. Fixes llvm#162902.

…eCost for bundled partial reductions

david-arm

LGTM!

sdesmalen-arm · 2025-10-21T13:55:21Z

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

    if (match(Widen, m_Sub(m_ZeroInt(), m_VPValue(Op)))) {
-      Widen = dyn_cast<VPWidenRecipe>(Op->getDefiningRecipe());
+      Widen = dyn_cast<VPWidenRecipe>(Op);
+      if (!Widen) {


I left a comment yesterday on @SamTebbs33' patch that took these changes: #147302 (comment)

Just to clarify what I meant, I don't think there's ever a case where this wouldn't be a WidenRecipe. For the case where Op would be a constant, then -Op would also have been a constant.

It probably makes sense to keep the assert though.

sdesmalen-arm · 2025-10-21T13:56:11Z

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

  } else if (isa<VPReductionPHIRecipe>(OpR)) {
    auto RedPhiOp1R = getOperand(1)->getDefiningRecipe();
-    if (isa<VPWidenCastRecipe>(RedPhiOp1R)) {
+    if (isa_and_nonnull<VPWidenCastRecipe>(RedPhiOp1R)) {


I left a comment yesterday on @SamTebbs33' patch that took these changes: #147302 (comment)

I realise I should probably have been a bit clearer here; I don't believe that OpR (which is the chain operand) can ever be nullptr. In contrast, getVecOp() (==getOperand(1)) can be so that should be special cased here, but it may be clearer to write that as:

// The partial.reduce.add's vector operand might not be a live-in. if (auto *RedPhiOp1R = getVecOp()->getDefiningRecipe()) { if (isa<VPWidenCastRecipe>(RedPhiOp1R) { ... } else if (auto Widen = dyn_cast<VPWidenRecipe>(RedPhiOp1R)) ... }

sdesmalen-arm · 2025-10-21T13:56:33Z

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

    InputTypeB = Ctx.Types.inferScalarType(ExtBR ? ExtBR->getOperand(0)
                                                 : Widen->getOperand(1));
    ExtAType = GetExtendKind(ExtAR);
    ExtBType = GetExtendKind(ExtBR);


I left a comment yesterday on @SamTebbs33' patch that took these changes: #147302 (comment)

(Note that this comment relates only to #147302 because that case couldn't be tested without the changes added in that PR, so there's nothing to do for that case in this PR)

…eCost for bundled partial reductions

fhahn requested review from SamTebbs33, david-arm, sdesmalen-arm and sjoerdmeijer October 13, 2025 10:52

llvmbot added vectorizers llvm:transforms labels Oct 13, 2025

sjoerdmeijer approved these changes Oct 13, 2025

View reviewed changes

SamTebbs33 reviewed Oct 13, 2025

View reviewed changes

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp Outdated Show resolved Hide resolved

david-arm reviewed Oct 14, 2025

View reviewed changes

SamTebbs33 reviewed Oct 14, 2025

View reviewed changes

SamTebbs33 mentioned this pull request Oct 14, 2025

[LV] Bundle partial reductions inside VPExpressionRecipe #147302

Merged

fhahn added 2 commits October 19, 2025 19:35

[VPlan] Handle live-in extend operands in partial reduction ::compute…

4435fa1

…Cost In some cases, extend operands can be live-ins. Use dyn_cast_if_present, to handle this case without crashing. Fixes llvm#162902.

!fixup add assertions, use dyn_cast_if_present

26fa3af

fhahn force-pushed the vplan-partial-cost-crash branch from d2532f8 to 26fa3af Compare October 19, 2025 18:59

SamTebbs33 added a commit to SamTebbs33/llvm-project that referenced this pull request Oct 20, 2025

Remove changes from llvm#163175 and add a shorter code path in comput…

db46c3b

…eCost for bundled partial reductions

david-arm approved these changes Oct 21, 2025

View reviewed changes

sdesmalen-arm reviewed Oct 21, 2025

View reviewed changes

SamTebbs33 mentioned this pull request Oct 23, 2025

[LV] Bundle (partial) reductions with a mul of a constant #162503

Open

SamTebbs33 added a commit to SamTebbs33/llvm-project that referenced this pull request Oct 23, 2025

Remove changes from llvm#163175 and add a shorter code path in comput…

0785325

…eCost for bundled partial reductions

[VPlan] Handle live-in extend operands in partial reduction ::computeCost #163175

Are you sure you want to change the base?

[VPlan] Handle live-in extend operands in partial reduction ::computeCost #163175

Conversation

fhahn commented Oct 13, 2025

Uh oh!

llvmbot commented Oct 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sjoerdmeijer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SamTebbs33 Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

david-arm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

llvmbot commented Oct 13, 2025 •

edited

Loading

SamTebbs33 Oct 14, 2025 •

edited

Loading