[llvm] [VPlan] Support live-ins without underlying IR in type analysis. (PR #80723)

[Florian Hahn via llvm-commits]

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

A VPlan contains multiple live-ins without underlying IR, like VFxUF or VectorTripCount. Trying to infer the scalar type of those causes a crash at the moment.

Update VPTypeAnalysis to take a VPlan in its constructor and assign types to those live-ins up front. All those live-ins share the type of the canonical IV.

>From 15b2c7c33347a9876093d8fb56451b55bf3e043c Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 5 Feb 2024 18:40:33 +0000
Subject: [PATCH] [VPlan] Support live-ins without underlying IR in type
 analysis.

A VPlan contains multiple live-ins without underlying IR, like VFxUF or
VectorTripCount. Trying to infer the scalar type of those causes a crash
at the moment.

Update VPTypeAnalysis to take a VPlan in its constructor and assign
types to those live-ins up front. All those live-ins share the type of
the canonical IV.
---
 llvm/lib/Transforms/Vectorize/VPlan.h           |  5 ++++-
 llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp | 17 +++++++++++------
 llvm/lib/Transforms/Vectorize/VPlanAnalysis.h   |  3 ++-
 .../Transforms/Vectorize/VPlanTransforms.cpp    |  8 ++++----
 4 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 162a3c4b195e5..43f10c315f2e9 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -238,7 +238,7 @@ struct VPTransformState {
                    DominatorTree *DT, IRBuilderBase &Builder,
                    InnerLoopVectorizer *ILV, VPlan *Plan, LLVMContext &Ctx)
       : VF(VF), UF(UF), LI(LI), DT(DT), Builder(Builder), ILV(ILV), Plan(Plan),
-        LVer(nullptr), TypeAnalysis(Ctx) {}
+        LVer(nullptr), TypeAnalysis(*Plan, Ctx) {}
 
   /// The chosen Vectorization and Unroll Factors of the loop being vectorized.
   ElementCount VF;
@@ -2932,6 +2932,9 @@ class VPlan {
     return BackedgeTakenCount;
   }
 
+  /// Return the backedge taken count of the original loop, if set.
+  VPValue *getBackedgeTakenCount() { return BackedgeTakenCount; }
+
   /// The vector trip count.
   VPValue &getVectorTripCount() { return VectorTripCount; }
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index b9ffe7e5b7af7..b71bac1546332 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -35,12 +35,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     CachedTypes[OtherV] = ResTy;
     return ResTy;
   }
-  case Instruction::ICmp: {
-    // TODO: Check if types for both operands agree. This also requires
-    // type-inference for the vector-trip-count, which is missing at the moment.
-    Type *ResTy = inferScalarType(R->getOperand(0));
-    return ResTy;
-  }
+  case Instruction::ICmp:
   case VPInstruction::FirstOrderRecurrenceSplice: {
     Type *ResTy = inferScalarType(R->getOperand(0));
     VPValue *OtherV = R->getOperand(1);
@@ -203,6 +198,16 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPReplicateRecipe *R) {
   llvm_unreachable("Unhandled opcode");
 }
 
+VPTypeAnalysis::VPTypeAnalysis(VPlan &Plan, LLVMContext &Ctx) : Ctx(Ctx) {
+  auto *CanIV = Plan.getCanonicalIV();
+  Type *CanIVTy = inferScalarType(CanIV);
+  CachedTypes[&Plan.getVectorTripCount()] = CanIVTy;
+  CachedTypes[&Plan.getVFxUF()] = CanIVTy;
+  CachedTypes[Plan.getTripCount()] = CanIVTy;
+  if (auto *BTC = Plan.getBackedgeTakenCount())
+    CachedTypes[BTC] = CanIVTy;
+}
+
 Type *VPTypeAnalysis::inferScalarType(const VPValue *V) {
   if (Type *CachedTy = CachedTypes.lookup(V))
     return CachedTy;
diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.h b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.h
index 7276641551ae8..6b218a39f605c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.h
@@ -15,6 +15,7 @@ namespace llvm {
 
 class LLVMContext;
 class VPValue;
+class VPlan;
 class VPBlendRecipe;
 class VPInstruction;
 class VPWidenRecipe;
@@ -47,7 +48,7 @@ class VPTypeAnalysis {
   Type *inferScalarTypeForRecipe(const VPReplicateRecipe *R);
 
 public:
-  VPTypeAnalysis(LLVMContext &Ctx) : Ctx(Ctx) {}
+  VPTypeAnalysis(VPlan &Plan, LLVMContext &Ctx);
 
   /// Infer the type of \p V. Returns the scalar type of \p V.
   Type *inferScalarType(const VPValue *V);
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 71f5285f90236..4f4febcd1403f 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -502,7 +502,7 @@ static VPValue *createScalarIVSteps(VPlan &Plan, const InductionDescriptor &ID,
   }
 
   // Truncate base induction if needed.
-  VPTypeAnalysis TypeInfo(SE.getContext());
+  VPTypeAnalysis TypeInfo(Plan, SE.getContext());
   Type *ResultTy = TypeInfo.inferScalarType(BaseIV);
   if (TruncI) {
     Type *TruncTy = TruncI->getType();
@@ -880,7 +880,7 @@ static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
 #ifndef NDEBUG
     // Verify that the cached type info is for both A and its users is still
     // accurate by comparing it to freshly computed types.
-    VPTypeAnalysis TypeInfo2(TypeInfo.getContext());
+    VPTypeAnalysis TypeInfo2(*R.getParent()->getPlan(), TypeInfo.getContext());
     assert(TypeInfo.inferScalarType(A) == TypeInfo2.inferScalarType(A));
     for (VPUser *U : A->users()) {
       auto *R = dyn_cast<VPRecipeBase>(U);
@@ -901,7 +901,7 @@ static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
 static void simplifyRecipes(VPlan &Plan, LLVMContext &Ctx) {
   ReversePostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>> RPOT(
       Plan.getEntry());
-  VPTypeAnalysis TypeInfo(Ctx);
+  VPTypeAnalysis TypeInfo(Plan, Ctx);
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
     for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
       simplifyRecipe(R, TypeInfo);
@@ -922,7 +922,7 @@ void VPlanTransforms::truncateToMinimalBitwidths(
   // other uses have different types for their operands, making them invalidly
   // typed.
   DenseMap<VPValue *, VPWidenCastRecipe *> ProcessedTruncs;
-  VPTypeAnalysis TypeInfo(Ctx);
+  VPTypeAnalysis TypeInfo(Plan, Ctx);
   VPBasicBlock *PH = Plan.getEntry();
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_deep(Plan.getVectorLoopRegion()))) {