[llvm] [VPlan] Add VPInst::getNumOperandsForOpcode, use to verify in ctor (NFC) (PR #142284)

[Florian Hahn via llvm-commits]

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

Add a new getNumOperandsForOpcode helper to determine the number of operands from the opcode. For now, it is used to verify the number operands at VPInstruction construction.

It returns -1 for a few opcodes where the number of operands cannot be determined (GEP, Switch, PHI, Call).

This can also be used in a follow-up to determine if a VPInstruction is masked based on the number of arguments.

>From 2ab3df5245bf142386eaed027688467b07c3a8f5 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 31 May 2025 16:04:10 +0100
Subject: [PATCH] [VPlan] Add VPInst::getNumOperandsForOpcode, use to verify in
 ctor (NFC)

Add a new getNumOperandsForOpcode helper to determine the number of
operands from the opcode. For now, it is used to verify the number
operands at VPInstruction construction.

It returns -1 for a few opcodes where the number of operands cannot be
determined (GEP, Switch, PHI, Call).

This can also be used in a follow-up to determine if a VPInstruction is
masked based on the number of arguments.
---
 llvm/lib/Transforms/Vectorize/VPlan.h         |  7 +++
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp | 59 ++++++++++++++++++-
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 44f0b6d964a6e..bd6a3247abd7d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -967,6 +967,13 @@ class VPInstruction : public VPRecipeWithIRFlags,
   /// value for lane \p Lane.
   Value *generatePerLane(VPTransformState &State, const VPLane &Lane);
 
+#if !defined(NDEBUG)
+  /// Return the number of operands determined by the opcode of the
+  /// VPInstruction. Returns -1 if the number of operands cannot be determined
+  /// directly by the opcode.
+  unsigned getNumOperandsForOpcode() const;
+#endif
+
 public:
   VPInstruction(unsigned Opcode, ArrayRef<VPValue *> Operands, DebugLoc DL = {},
                 const Twine &Name = "")
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index a4831ea7c11f7..5c6e4aeaf3cad 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -413,8 +413,62 @@ VPInstruction::VPInstruction(unsigned Opcode, ArrayRef<VPValue *> Operands,
       Opcode(Opcode), Name(Name.str()) {
   assert(flagsValidForOpcode(getOpcode()) &&
          "Set flags not supported for the provided opcode");
+  assert((getNumOperandsForOpcode() == -1u ||
+          getNumOperandsForOpcode() == getNumOperands()) &&
+         "number of operands does not match opcode");
 }
 
+#ifndef NDEBUG
+unsigned VPInstruction::getNumOperandsForOpcode() const {
+  if (Instruction::isUnaryOp(getOpcode()) || Instruction::isCast(getOpcode()))
+    return 1;
+
+  if (Instruction::isBinaryOp(getOpcode()))
+    return 2;
+
+  switch (getOpcode()) {
+  case VPInstruction::StepVector:
+    return 0;
+  case Instruction::Alloca:
+  case Instruction::ExtractValue:
+  case Instruction::Freeze:
+  case Instruction::Load:
+  case VPInstruction::AnyOf:
+  case VPInstruction::BranchOnCond:
+  case VPInstruction::CalculateTripCountMinusVF:
+  case VPInstruction::CanonicalIVIncrementForPart:
+  case VPInstruction::ExplicitVectorLength:
+  case VPInstruction::ExtractLastElement:
+  case VPInstruction::ExtractPenultimateElement:
+  case VPInstruction::FirstActiveLane:
+  case VPInstruction::Not:
+    return 1;
+
+  case Instruction::ICmp:
+  case Instruction::FCmp:
+  case Instruction::Store:
+  case VPInstruction::ActiveLaneMask:
+  case VPInstruction::BranchOnCount:
+  case VPInstruction::ComputeReductionResult:
+  case VPInstruction::FirstOrderRecurrenceSplice:
+  case VPInstruction::LogicalAnd:
+  case VPInstruction::WideIVStep:
+  case VPInstruction::PtrAdd:
+    return 2;
+  case Instruction::Select:
+  case VPInstruction::ComputeFindLastIVResult:
+    return 3;
+  case Instruction::Call:
+  case Instruction::PHI:
+  case Instruction::GetElementPtr:
+  case Instruction::Switch:
+    // Cannot determine the number of operands from the opcode.
+    return -1u;
+  }
+  llvm_unreachable("all cases should be handled above");
+}
+#endif
+
 bool VPInstruction::doesGeneratePerAllLanes() const {
   return Opcode == VPInstruction::PtrAdd && !vputils::onlyFirstLaneUsed(this);
 }
@@ -2706,7 +2760,10 @@ static void scalarizeInstruction(const Instruction *Instr,
 
   // Replace the operands of the cloned instructions with their scalar
   // equivalents in the new loop.
-  for (const auto &I : enumerate(RepRecipe->operands())) {
+  auto OpRange = RepRecipe->operands();
+  if (isa<CallBase>(Cloned))
+    OpRange = drop_end(OpRange);
+  for (const auto &I : enumerate(OpRange)) {
     auto InputLane = Lane;
     VPValue *Operand = I.value();
     if (vputils::isSingleScalar(Operand))