[llvm] [VPlan] Factor first-lane-scalar logic in VPI::execute (NFC) (PR #196207)

[Ramkumar Ramachandra via llvm-commits]

https://github.com/artagnon updated https://github.com/llvm/llvm-project/pull/196207

>From 27aa625bbf7c23108a5ba3dc66df82e85a6fe91f Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <artagnon at tenstorrent.com>
Date: Thu, 7 May 2026 00:07:53 +0100
Subject: [PATCH] [VPlan] Factor first-lane-scalar logic in VPI::execute (NFC)

---
 llvm/lib/Transforms/Vectorize/VPlan.h         |  8 +--
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp | 50 ++++++++-----------
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index daabca902724b..fe87545fec4d6 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1448,10 +1448,10 @@ class LLVM_ABI_FOR_TEST VPInstruction : public VPRecipeWithIRFlags,
   /// needed.
   bool canGenerateScalarForFirstLane() const;
 
-  /// Utility methods serving execute(): generates a single vector instance of
-  /// the modeled instruction. \returns the generated value. . In some cases an
-  /// existing value is returned rather than a generated one.
-  Value *generate(VPTransformState &State);
+  /// Utility method serving execute: Generates either a single-scalar or vector
+  /// value. \p IsSingleScalar determines whether to generate a single-scalar
+  /// value.
+  Value *generate(VPTransformState &State, bool IsSingleScalar);
 
   /// Returns true if the VPInstruction does not need masking.
   bool alwaysUnmasked() const {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 452589bafd534..c5cbc3cd45cab 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -742,13 +742,12 @@ static Instruction::BinaryOps getSubRecurOpcode(RecurKind Kind) {
   llvm_unreachable("RecurKind should be Sub/FSub.");
 }
 
-Value *VPInstruction::generate(VPTransformState &State) {
+Value *VPInstruction::generate(VPTransformState &State, bool IsSingleScalar) {
   IRBuilderBase &Builder = State.Builder;
 
   if (Instruction::isBinaryOp(getOpcode())) {
-    bool OnlyFirstLaneUsed = vputils::onlyFirstLaneUsed(this);
-    Value *A = State.get(getOperand(0), OnlyFirstLaneUsed);
-    Value *B = State.get(getOperand(1), OnlyFirstLaneUsed);
+    Value *A = State.get(getOperand(0), IsSingleScalar);
+    Value *B = State.get(getOperand(1), IsSingleScalar);
     auto *Res =
         Builder.CreateBinOp((Instruction::BinaryOps)getOpcode(), A, B, Name);
     if (auto *I = dyn_cast<Instruction>(Res))
@@ -768,8 +767,7 @@ Value *VPInstruction::generate(VPTransformState &State) {
 
   switch (getOpcode()) {
   case VPInstruction::Not: {
-    bool OnlyFirstLaneUsed = vputils::onlyFirstLaneUsed(this);
-    Value *A = State.get(getOperand(0), OnlyFirstLaneUsed);
+    Value *A = State.get(getOperand(0), IsSingleScalar);
     return Builder.CreateNot(A, Name);
   }
   case Instruction::ExtractElement: {
@@ -788,26 +786,24 @@ Value *VPInstruction::generate(VPTransformState &State) {
     return Builder.CreateInsertElement(Vec, Elt, Idx, Name);
   }
   case Instruction::Freeze: {
-    Value *Op = State.get(getOperand(0), vputils::onlyFirstLaneUsed(this));
+    Value *Op = State.get(getOperand(0), IsSingleScalar);
     return Builder.CreateFreeze(Op, Name);
   }
   case Instruction::FCmp:
   case Instruction::ICmp: {
-    bool OnlyFirstLaneUsed = vputils::onlyFirstLaneUsed(this);
-    Value *A = State.get(getOperand(0), OnlyFirstLaneUsed);
-    Value *B = State.get(getOperand(1), OnlyFirstLaneUsed);
+    Value *A = State.get(getOperand(0), IsSingleScalar);
+    Value *B = State.get(getOperand(1), IsSingleScalar);
     return Builder.CreateCmp(getPredicate(), A, B, Name);
   }
   case Instruction::PHI: {
     llvm_unreachable("should be handled by VPPhi::execute");
   }
   case Instruction::Select: {
-    bool OnlyFirstLaneUsed = vputils::onlyFirstLaneUsed(this);
     Value *Cond =
         State.get(getOperand(0),
-                  OnlyFirstLaneUsed || vputils::isSingleScalar(getOperand(0)));
-    Value *Op1 = State.get(getOperand(1), OnlyFirstLaneUsed);
-    Value *Op2 = State.get(getOperand(2), OnlyFirstLaneUsed);
+                  IsSingleScalar || vputils::isSingleScalar(getOperand(0)));
+    Value *Op1 = State.get(getOperand(1), IsSingleScalar);
+    Value *Op2 = State.get(getOperand(2), IsSingleScalar);
     return Builder.CreateSelectFMF(Cond, Op1, Op2, getFastMathFlagsOrNone(),
                                    Name);
   }
@@ -1023,13 +1019,13 @@ Value *VPInstruction::generate(VPTransformState &State) {
     return Builder.CreateLogicalOr(A, B, Name);
   }
   case VPInstruction::PtrAdd: {
-    assert((State.VF.isScalar() || vputils::onlyFirstLaneUsed(this)) &&
-           "can only generate first lane for PtrAdd");
-    Value *Ptr = State.get(getOperand(0), VPLane(0));
-    Value *Addend = State.get(getOperand(1), VPLane(0));
+    assert(IsSingleScalar && "Can only generate first lane for PtrAdd");
+    Value *Ptr = State.get(getOperand(0), IsSingleScalar);
+    Value *Addend = State.get(getOperand(1), IsSingleScalar);
     return Builder.CreatePtrAdd(Ptr, Addend, Name, getGEPNoWrapFlags());
   }
   case VPInstruction::WidePtrAdd: {
+    assert(!IsSingleScalar && "Cannot generate scalar value for WidePtrAdd");
     Value *Ptr =
         State.get(getOperand(0), vputils::isSingleScalar(getOperand(0)));
     Value *Addend = State.get(getOperand(1));
@@ -1640,20 +1636,18 @@ void VPInstruction::execute(VPTransformState &State) {
   assert(hasRequiredFlagsForOpcode(getOpcode(), getScalarType()) &&
          "Opcode requires specific flags to be set");
   State.Builder.setFastMathFlags(getFastMathFlagsOrNone());
-  Value *GeneratedValue = generate(State);
+  bool GenerateScalar =
+      State.VF.isScalar() || (canGenerateScalarForFirstLane() &&
+                              (vputils::onlyFirstLaneUsed(this) ||
+                               isVectorToScalar() || isSingleScalar()));
+  Value *GeneratedValue = generate(State, GenerateScalar);
   if (!hasResult())
     return;
   assert(GeneratedValue && "generate must produce a value");
-  bool GeneratesPerFirstLaneOnly = canGenerateScalarForFirstLane() &&
-                                   (vputils::onlyFirstLaneUsed(this) ||
-                                    isVectorToScalar() || isSingleScalar());
-  assert((((GeneratedValue->getType()->isVectorTy() ||
-            GeneratedValue->getType()->isStructTy()) ==
-           !GeneratesPerFirstLaneOnly) ||
-          State.VF.isScalar()) &&
+  assert(((GeneratedValue->getType()->isVectorTy() ||
+           GeneratedValue->getType()->isStructTy()) == !GenerateScalar) &&
          "scalar value but not only first lane defined");
-  State.set(this, GeneratedValue,
-            /*IsScalar*/ GeneratesPerFirstLaneOnly);
+  State.set(this, GeneratedValue, GenerateScalar);
   if (getOpcode() == VPInstruction::ResumeForEpilogue ||
       getOpcode() == Instruction::Freeze) {
     // FIXME: This is a workaround to enable reliable updates of the scalar loop