[VPlan] Use VPTypeInfo in simplifyRecipes.

Replace getTypeForVPValue with the recently added, more general
VPTypeAnalysis.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -181,6 +181,11 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPReplicateRecipe *R) {
     return inferScalarType(R->getOperand(0));
   case Instruction::Load:
     return cast<LoadInst>(R->getUnderlyingInstr())->getType();
+  case Instruction::Store:
+    // FIXME: VPReplicateRecipes with store opcodes still define a result
+    // VPValue, so we need to handle them here. Remove the code here once this
+    // is modeled accurately in VPlan.
+    return Type::getVoidTy(Ctx);
   default:
     break;
   }

llvm/lib/Transforms/Vectorize/VPlanAnalysis.h

@@ -54,6 +54,9 @@ class VPTypeAnalysis {
 
   /// Infer the type of \p V. Returns the scalar type of \p V.
   Type *inferScalarType(const VPValue *V);
+
+  /// Return the LLVMContext used by the analysis.
+  LLVMContext &getContext() { return Ctx; }
 };
 
 } // end namespace llvm

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -13,6 +13,7 @@
 
 #include "VPlanTransforms.h"
 #include "VPRecipeBuilder.h"
+#include "VPlanAnalysis.h"
 #include "VPlanCFG.h"
 #include "VPlanDominatorTree.h"
 #include "llvm/ADT/PostOrderIterator.h"
@@ -802,17 +803,8 @@ static unsigned getOpcodeForRecipe(VPRecipeBase &R) {
   return 0;
 }
 
-/// Return the scalar size in bits for \p VPV if possible.
-static Type *getTypeForVPValue(VPValue *VPV) {
-  // TODO: Replace with VPlan type inference once ready.
-  if (auto *VPC = dyn_cast<VPWidenCastRecipe>(VPV))
-    return VPC->getResultType();
-  auto *UV = VPV->getUnderlyingValue();
-  return UV ? UV->getType() : nullptr;
-}
-
 /// Try to simplify recipe \p R.
-static void simplifyRecipe(VPRecipeBase &R) {
+static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
   switch (getOpcodeForRecipe(R)) {
   case Instruction::Mul: {
     VPValue *A = R.getOperand(0);
@@ -829,9 +821,23 @@ static void simplifyRecipe(VPRecipeBase &R) {
       break;
     VPValue *A = Zext->getOperand(0);
     VPValue *Trunc = R.getVPSingleValue();
-    Type *TruncToTy = getTypeForVPValue(Trunc);
-    if (TruncToTy && TruncToTy == getTypeForVPValue(A))
+    Type *TruncToTy = TypeInfo.inferScalarType(Trunc);
+    if (TruncToTy && TruncToTy == TypeInfo.inferScalarType(A))
       Trunc->replaceAllUsesWith(A);
+
+#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());
+    assert(TypeInfo.inferScalarType(A) == TypeInfo2.inferScalarType(A));
+    for (VPUser *U : A->users()) {
+      auto *R = dyn_cast<VPRecipeBase>(U);
+      if (!R)
+        continue;
+      for (VPValue *VPV : R->definedValues())
+        assert(TypeInfo.inferScalarType(VPV) == TypeInfo2.inferScalarType(VPV));
+    }
+#endif
     break;
   }
   default:
@@ -840,12 +846,13 @@ static void simplifyRecipe(VPRecipeBase &R) {
 }
 
 /// Try to simplify the recipes in \p Plan.
-static void simplifyRecipes(VPlan &Plan) {
+static void simplifyRecipes(VPlan &Plan, LLVMContext &Ctx) {
   ReversePostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>> RPOT(
       Plan.getEntry());
+  VPTypeAnalysis TypeInfo(Ctx);
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
     for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
-      simplifyRecipe(R);
+      simplifyRecipe(R, TypeInfo);
     }
   }
 }
@@ -855,7 +862,7 @@ void VPlanTransforms::optimize(VPlan &Plan, ScalarEvolution &SE) {
   removeRedundantInductionCasts(Plan);
 
   optimizeInductions(Plan, SE);
-  simplifyRecipes(Plan);
+  simplifyRecipes(Plan, SE.getContext());
   removeDeadRecipes(Plan);
 
   createAndOptimizeReplicateRegions(Plan);

llvm/test/Transforms/LoopVectorize/cast-induction.ll

@@ -53,12 +53,8 @@ define void @redundant_iv_cast(ptr %dst) {
 ; IC2-NEXT:  [[OFFSET_IDX:%.+]] = trunc i32 [[CAN_IV]] to i16
 ; IC2-NEXT:  [[P0:%.+]] = add i16 [[OFFSET_IDX]], 0
 ; IC2-NEXT:  [[P1:%.+]] = add i16 [[OFFSET_IDX]], 1
-; IC2-NEXT:  [[Z0:%.+]] = zext i16 [[P0]] to i32
-; IC2-NEXT:  [[Z1:%.+]] = zext i16 [[P1]] to i32
-; IC2-NEXT:  [[T0:%.+]] = trunc i32 [[Z0]] to i16
-; IC2-NEXT:  [[T1:%.+]] = trunc i32 [[Z1]] to i16
-; IC2:       store i16 [[T0]]
-; IC2-NEXT:  store i16 [[T1]]
+; IC2:       store i16 [[P0]]
+; IC2-NEXT:  store i16 [[P1]]
 ;
 entry:
   br label %loop