[llvm] [VectorUtils][VPlan] Consolidate VPWidenIntrinsicRecipe::onlyFirstLaneUsed and isVectorIntrinsicWithScalarOpAtArg (PR #137497)

[Luke Lau via llvm-commits]

https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/137497

We can reuse isVectorIntrinsicWithScalarOpAtArg in VectorUtils to determine if only the first lane will be used for a VPWidenIntrinsicRecipe, provided that we also move the VP EVL operand check into it.

This was needed by a local patch I was working on that created a VPWidenIntrinsicRecipe with a VP intrinsic, and prevents the need to update the scalar arguments in two places.


>From a3475ac997ec6b20cce2225093a76d1ed7f22e9a Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Sun, 27 Apr 2025 13:42:16 +0800
Subject: [PATCH] [VectorUtils][VPlan] Consolidate
 VPWidenIntrinsicRecipe::onlyFirstLaneUsed and
 isVectorIntrinsicWithScalarOpAtArg

We can reuse isVectorIntrinsicWithScalarOpAtArg in VectorUtils to determine if only the first lane will be used for a VPWidenIntrinsicRecipe, provided that we also move the VP EVL operand check into it.

This was needed by a local patch I was working on that created a VPWidenIntrinsicRecipe with a VP intrinsic, and prevents the need to update the scalar arguments in two places.
---
 llvm/lib/Analysis/VectorUtils.cpp              | 7 ++++++-
 llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp | 6 ++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp
index 6448c372f5d5d..b614ac50d2f82 100644
--- a/llvm/lib/Analysis/VectorUtils.cpp
+++ b/llvm/lib/Analysis/VectorUtils.cpp
@@ -149,6 +149,11 @@ bool llvm::isVectorIntrinsicWithScalarOpAtArg(Intrinsic::ID ID,
   if (TTI && Intrinsic::isTargetIntrinsic(ID))
     return TTI->isTargetIntrinsicWithScalarOpAtArg(ID, ScalarOpdIdx);
 
+  // Vector predication intrinsics only demand the the first lane the last
+  // operand (the EVL operand).
+  if (VPIntrinsic::getVectorLengthParamPos(ID) == ScalarOpdIdx)
+    return true;
+
   switch (ID) {
   case Intrinsic::abs:
   case Intrinsic::vp_abs:
@@ -166,7 +171,7 @@ bool llvm::isVectorIntrinsicWithScalarOpAtArg(Intrinsic::ID ID,
   case Intrinsic::umul_fix_sat:
     return (ScalarOpdIdx == 2);
   case Intrinsic::experimental_vp_splice:
-    return ScalarOpdIdx == 2 || ScalarOpdIdx == 4 || ScalarOpdIdx == 5;
+    return ScalarOpdIdx == 2 || ScalarOpdIdx == 4;
   default:
     return false;
   }
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 8ec092ce9a905..787b416ee9248 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -1373,10 +1373,8 @@ StringRef VPWidenIntrinsicRecipe::getIntrinsicName() const {
 
 bool VPWidenIntrinsicRecipe::onlyFirstLaneUsed(const VPValue *Op) const {
   assert(is_contained(operands(), Op) && "Op must be an operand of the recipe");
-  // Vector predication intrinsics only demand the the first lane the last
-  // operand (the EVL operand).
-  return VPIntrinsic::isVPIntrinsic(VectorIntrinsicID) &&
-         Op == getOperand(getNumOperands() - 1);
+  unsigned Idx = std::distance(op_begin(), find(operands(), Op));
+  return isVectorIntrinsicWithScalarOpAtArg(VectorIntrinsicID, Idx, nullptr);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)