[llvm] [VPlan] Compute cost of intrinsics directly for VPReplicateRecipe (NFCI). (PR #154617)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/154617

>From bddd21b13eb007ac2d51f76a43638ef559b4311c Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 18 Aug 2025 19:04:27 +0100
Subject: [PATCH 1/2] [VPlan] Compute cost of intrinsics directly for
 VPReplicateRecipe (NFCI).

Handle intrinsic calls in VPReplicateRecipe::computeCost. There are some
intrinsics pseudo intrinsics for which the computed cost is known zero,
so we handle those up front.

Depends on https://github.com/llvm/llvm-project/pull/154291.
---
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp | 37 +++++++++++++++----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 2754a389ea12d..ba17773527e63 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -3063,17 +3063,41 @@ InstructionCost VPReplicateRecipe::computeCost(ElementCount VF,
   case Instruction::Call: {
     auto *CalledFn =
         cast<Function>(getOperand(getNumOperands() - 1)->getLiveInIRValue());
-    if (CalledFn->isIntrinsic())
-      break;
 
+    SmallVector<const VPValue *> ArgOps(drop_end(operands()));
     SmallVector<Type *, 4> Tys;
-    for (VPValue *ArgOp : drop_end(operands()))
+    for (const VPValue *ArgOp : ArgOps)
       Tys.push_back(Ctx.Types.inferScalarType(ArgOp));
+
+    if (CalledFn->isIntrinsic())
+      // Various pseudo-intrinsics with costs of 0 are scalarized instead of
+      // vectorized via VPWidenIntrinsicRecipe. Return 0 for them early.
+      switch (CalledFn->getIntrinsicID()) {
+      case Intrinsic::assume:
+      case Intrinsic::lifetime_end:
+      case Intrinsic::lifetime_start:
+      case Intrinsic::sideeffect:
+      case Intrinsic::pseudoprobe:
+      case Intrinsic::experimental_noalias_scope_decl: {
+        assert(getCostForIntrinsics(CalledFn->getIntrinsicID(), ArgOps, *this,
+                                    ElementCount::getFixed(1), Ctx) == 0 && "scalarizing intrinsic should be free");
+        return InstructionCost(0);
+      }
+      default:
+        break;
+      }
+
     Type *ResultTy = Ctx.Types.inferScalarType(this);
     InstructionCost ScalarCallCost =
         Ctx.TTI.getCallInstrCost(CalledFn, ResultTy, Tys, Ctx.CostKind);
-    if (isSingleScalar())
+    if (isSingleScalar()) {
+      if (CalledFn->isIntrinsic())
+        ScalarCallCost = std::min(
+            ScalarCallCost,
+            getCostForIntrinsics(CalledFn->getIntrinsicID(), ArgOps, *this,
+                                 ElementCount::getFixed(1), Ctx));
       return ScalarCallCost;
+    }
 
     if (VF.isScalable())
       return InstructionCost::getInvalid();
@@ -3094,7 +3118,7 @@ InstructionCost VPReplicateRecipe::computeCost(ElementCount VF,
       // incur any overhead.
       SmallPtrSet<const VPValue *, 4> UniqueOperands;
       Tys.clear();
-      for (auto *Op : drop_end(operands())) {
+      for (auto *Op : ArgOps) {
         if (Op->isLiveIn() || isa<VPReplicateRecipe, VPPredInstPHIRecipe>(Op) ||
             !UniqueOperands.insert(Op).second)
           continue;
@@ -3104,8 +3128,7 @@ InstructionCost VPReplicateRecipe::computeCost(ElementCount VF,
           Ctx.TTI.getOperandsScalarizationOverhead(Tys, Ctx.CostKind);
     }
 
-    return ScalarCallCost * (isSingleScalar() ? 1 : VF.getFixedValue()) +
-           ScalarizationCost;
+    return ScalarCallCost * VF.getFixedValue() + ScalarizationCost;
   }
   case Instruction::Add:
   case Instruction::Sub:

>From 36dfca06172c49c3cb0f855c6db09e17efb380d1 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 27 Aug 2025 20:21:33 +0100
Subject: [PATCH 2/2] !fixup fix formatting

---
 llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index ba17773527e63..bd9a93ed57b8a 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -3080,7 +3080,8 @@ InstructionCost VPReplicateRecipe::computeCost(ElementCount VF,
       case Intrinsic::pseudoprobe:
       case Intrinsic::experimental_noalias_scope_decl: {
         assert(getCostForIntrinsics(CalledFn->getIntrinsicID(), ArgOps, *this,
-                                    ElementCount::getFixed(1), Ctx) == 0 && "scalarizing intrinsic should be free");
+                                    ElementCount::getFixed(1), Ctx) == 0 &&
+               "scalarizing intrinsic should be free");
         return InstructionCost(0);
       }
       default: