[llvm] [VPlan] Don't create ExtractElement recipes for scalar plans. (PR #131604)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/131604

ExtractElements are no-ops for scalar VPlans. Don't introduce them in handleUncountableEarlyExit if the plan has only a scalar VF.

This fixes a crash trying to compute the cost of ExtractElement after 26ecf978951b79.

>From 4134f414da221d734674b15fbe65e28e178bb94b Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 17 Mar 2025 11:12:47 +0000
Subject: [PATCH] [VPlan] Don't create ExtractElement recipes for scalar plans.

ExtractElements are no-ops for scalar VPlans. Don't introduce them in
handleUncountableEarlyExit if the plan has only a scalar VF.

This fixes a crash trying to compute the cost of ExtractElement after
26ecf978951b79.
---
 .../Transforms/Vectorize/VPlanTransforms.cpp  |  2 +-
 .../LoopVectorize/AArch64/early_exit_costs.ll | 29 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 9aae383d35d91..ad1006ceb2069 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2158,7 +2158,7 @@ void VPlanTransforms::handleUncountableEarlyExit(
       ExitIRI->extractLastLaneOfOperand(MiddleBuilder);
     }
     // Add the incoming value from the early exit.
-    if (!IncomingFromEarlyExit->isLiveIn()) {
+    if (!IncomingFromEarlyExit->isLiveIn() && !Plan.hasScalarVFOnly()) {
       VPValue *FirstActiveLane = EarlyExitB.createNaryOp(
           VPInstruction::FirstActiveLane, {EarlyExitTakenCond}, nullptr,
           "first.active.lane");
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll b/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll
index 4d7c5d088034d..6c794b7503df0 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/early_exit_costs.ll
@@ -1,4 +1,3 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
 ; REQUIRES: asserts
 ; RUN: opt -S < %s -p loop-vectorize -enable-early-exit-vectorization -disable-output \
 ; RUN:   -debug-only=loop-vectorize 2>&1 | FileCheck %s --check-prefixes=CHECK
@@ -87,4 +86,32 @@ loop.end:
   ret i64 %retval
 }
 
+define i64 @vectorization_not_profitable_due_to_trunc(ptr dereferenceable(800) %src) {
+; CHECK-LABEL: LV: Checking a loop in 'vectorization_not_profitable_due_to_trunc'
+; CHECK: LV: Selecting VF: 1.
+; CHECK-NEXT: Calculating cost of work in exit block vector.early.exit:
+; CHECK-NEXT: LV: Vectorization is possible but not beneficial.
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  %gep.src = getelementptr inbounds i64, ptr %src, i64 %iv
+  %l = load i64, ptr %gep.src, align 1
+  %t = trunc i64 %l to i1
+  br i1 %t, label %exit.0, label %loop.latch
+
+loop.latch:
+  %iv.next = add i64 %iv, 1
+  %ec = icmp eq i64 %iv.next, 100
+  br i1 %ec, label %exit.1, label %loop.header
+
+exit.0:
+  %res = phi i64 [ %l, %loop.header ]
+  ret i64 %res
+
+exit.1:
+  ret i64 0
+}
+
 attributes #1 = { "target-features"="+sve" vscale_range(1,16) }