[llvm] [VPlan] Replace first-lane uses of BuildVector. (PR #206566)

[Florian Hahn via llvm-commits]

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

Replace uses of the first lane of a BuildVector with the first operand.

>From 02f7450766a142a092e4a30c94a654a412dc85b6 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 29 Jun 2026 19:54:35 +0100
Subject: [PATCH] [VPlan] Replace first-lane uses of BuildVector.

Replace uses of the first lane of a BuildVector with the first operand.
---
 llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp  |  8 ++++++++
 .../VPlan/buildvector-first-lane-only.ll           | 14 +++++---------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 46c63bb47b165..93e006c5bcdf6 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1735,6 +1735,14 @@ static void simplifyRecipe(VPSingleDefRecipe *Def) {
     return;
   }
 
+  // Replace uses of a BuildVector by users that only use its first lane with
+  // its first operand directly.
+  if (match(Def, m_BuildVector())) {
+    Def->replaceUsesWithIf(
+        Def->getOperand(0),
+        [Def](VPUser &U, unsigned) { return U.usesFirstLaneOnly(Def); });
+  }
+
   // Look through broadcast of single-scalar when used as select conditions; in
   // that case the scalar condition can be used directly.
   if (match(Def,
diff --git a/llvm/test/Transforms/LoopVectorize/VPlan/buildvector-first-lane-only.ll b/llvm/test/Transforms/LoopVectorize/VPlan/buildvector-first-lane-only.ll
index 54614ad7c2035..8361c2d3a59b2 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/buildvector-first-lane-only.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/buildvector-first-lane-only.ll
@@ -19,19 +19,15 @@ define i32 @last_active_lane_live_out(i32 %x) {
 ; CHECK-NEXT:  Successor(s): middle.block
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  middle.block:
-; CHECK-NEXT:    vp<[[VP5:%[0-9]+]]> = SCALAR-STEPS ir<0>, ir<1>, ir<4>, ir<1>
-; CHECK-NEXT:    vp<[[VP6:%[0-9]+]]> = SCALAR-STEPS ir<0>, ir<1>, ir<4>, ir<2>
-; CHECK-NEXT:    vp<[[VP7:%[0-9]+]]> = SCALAR-STEPS ir<0>, ir<1>, ir<4>, ir<3>
-; CHECK-NEXT:    EMIT vp<[[VP8:%[0-9]+]]> = buildvector ir<0>, vp<[[VP5]]>, vp<[[VP6]]>, vp<[[VP7]]>
-; CHECK-NEXT:    EMIT vp<%active.lane.mask> = active lane mask vp<[[VP8]]>, ir<2>, ir<1>
-; CHECK-NEXT:    EMIT vp<[[VP9:%[0-9]+]]> = not vp<%active.lane.mask>
-; CHECK-NEXT:    EMIT vp<%first.inactive.lane> = first-active-lane vp<[[VP9]]>
+; CHECK-NEXT:    EMIT vp<%active.lane.mask> = active lane mask ir<0>, ir<2>, ir<1>
+; CHECK-NEXT:    EMIT vp<[[VP5:%[0-9]+]]> = not vp<%active.lane.mask>
+; CHECK-NEXT:    EMIT vp<%first.inactive.lane> = first-active-lane vp<[[VP5]]>
 ; CHECK-NEXT:    EMIT vp<%last.active.lane> = sub vp<%first.inactive.lane>, ir<1>
-; CHECK-NEXT:    EMIT vp<[[VP10:%[0-9]+]]> = extractelement vp<[[VP4]]>, vp<%last.active.lane>
+; CHECK-NEXT:    EMIT vp<[[VP6:%[0-9]+]]> = extractelement vp<[[VP4]]>, vp<%last.active.lane>
 ; CHECK-NEXT:  Successor(s): ir-bb<exit>
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  ir-bb<exit>:
-; CHECK-NEXT:    IR   %red.lcssa = phi i32 [ %red, %loop ] (extra operand: vp<[[VP10]]> from middle.block)
+; CHECK-NEXT:    IR   %red.lcssa = phi i32 [ %red, %loop ] (extra operand: vp<[[VP6]]> from middle.block)
 ; CHECK-NEXT:  No successors
 ; CHECK-NEXT:  }
 ;