[llvm] Reapply "[VPlan] Replace first-lane uses of BuildVector." (#206774) (PR #207512)

[Florian Hahn via llvm-commits]

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

This reverts commit 638d71403a8598bb061b874ab39b8f9c6f0a97f0.

Recommit with fix for root cause of the revert:
VPWidenRecipe::onlyFirstLaneUsed was incorrectly assuming all operands defined outside loop region are invariant, which is not true (e.g. BuildVector). Test added in e173245deeaa2340.

Original message:
Replace uses of the first lane of a BuildVector with the first operand.

>From 0606a82519b702d5329a8ee673f7304f28560875 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 4 Jul 2026 11:08:40 +0100
Subject: [PATCH] Reapply "[VPlan] Replace first-lane uses of BuildVector."
 (#206774)

This reverts commit 638d71403a8598bb061b874ab39b8f9c6f0a97f0.

Recommit with fix for root cause of the revert:
VPWidenRecipe::onlyFirstLaneUsed was incorrectly assuming all operands
defined outside loop region are invariant, which is not true (e.g.
BuildVector). Test added in e173245deeaa2340.

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

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 3e7be2d67733a..b2d179d8e6d98 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1868,7 +1868,7 @@ class LLVM_ABI_FOR_TEST VPWidenRecipe : public VPRecipeWithIRFlags,
     assert(is_contained(operands(), Op) &&
            "Op must be an operand of the recipe");
     return Opcode == Instruction::Select && Op == getOperand(0) &&
-           Op->isDefinedOutsideLoopRegions();
+           isa<VPIRValue>(Op);
   }
 };
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index e5d16ccbec434..bd260a0d12362 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1745,6 +1745,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 04f3aa3accbc8..297eccd66c2d5 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/buildvector-first-lane-only.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/buildvector-first-lane-only.ll
@@ -20,19 +20,15 @@ define i16 @last_active_lane_live_out(i32 %x) {
 ; CHECK-NEXT:  Successor(s): middle.block
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  middle.block:
-; CHECK-NEXT:    vp<[[VP6:%[0-9]+]]> = SCALAR-STEPS ir<0>, ir<1>, ir<4>, ir<1>
-; CHECK-NEXT:    vp<[[VP7:%[0-9]+]]> = SCALAR-STEPS ir<0>, ir<1>, ir<4>, ir<2>
-; CHECK-NEXT:    vp<[[VP8:%[0-9]+]]> = SCALAR-STEPS ir<0>, ir<1>, ir<4>, ir<3>
-; CHECK-NEXT:    EMIT vp<[[VP9:%[0-9]+]]> = buildvector ir<0>, vp<[[VP6]]>, vp<[[VP7]]>, vp<[[VP8]]>
-; CHECK-NEXT:    EMIT vp<%active.lane.mask> = active lane mask vp<[[VP9]]>, ir<2>, ir<1>
-; CHECK-NEXT:    EMIT vp<[[VP10:%[0-9]+]]> = not vp<%active.lane.mask>
-; CHECK-NEXT:    EMIT vp<%first.inactive.lane> = first-active-lane vp<[[VP10]]>
+; CHECK-NEXT:    EMIT vp<%active.lane.mask> = active lane mask ir<0>, ir<2>, ir<1>
+; CHECK-NEXT:    EMIT vp<[[VP6:%[0-9]+]]> = not vp<%active.lane.mask>
+; CHECK-NEXT:    EMIT vp<%first.inactive.lane> = first-active-lane vp<[[VP6]]>
 ; CHECK-NEXT:    EMIT vp<%last.active.lane> = sub vp<%first.inactive.lane>, ir<1>
-; CHECK-NEXT:    EMIT vp<[[VP11:%[0-9]+]]> = extractelement vp<[[VP5]]>, vp<%last.active.lane>
+; CHECK-NEXT:    EMIT vp<[[VP7:%[0-9]+]]> = extractelement vp<[[VP5]]>, vp<%last.active.lane>
 ; CHECK-NEXT:  Successor(s): ir-bb<exit>
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  ir-bb<exit>:
-; CHECK-NEXT:    IR   %t.lcssa = phi i16 [ %t, %loop ] (extra operand: vp<[[VP11]]> from middle.block)
+; CHECK-NEXT:    IR   %t.lcssa = phi i16 [ %t, %loop ] (extra operand: vp<[[VP7]]> from middle.block)
 ; CHECK-NEXT:  No successors
 ; CHECK-NEXT:  }
 ;