Remove erased type from VTypeAnalysis cache, move assertion into all simplifications

Author: lukel97

llvm/lib/Transforms/Vectorize/VPlanAnalysis.h

@@ -63,6 +63,9 @@ class VPTypeAnalysis {
 
   /// Return the LLVMContext used by the analysis.
   LLVMContext &getContext() { return Ctx; }
+
+  /// Remove \p V from the cache. You must call this after a value is erased.
+  void erase(VPValue *V) { CachedTypes.erase(V); }
 };
 
 // Collect a VPlan's ephemeral recipes (those used only by an assume).

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -966,18 +966,6 @@ static VPValue *simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
         return VPC;
       }
     }
-#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(
-        R.getParent()->getPlan()->getCanonicalIV()->getScalarType());
-    assert(TypeInfo.inferScalarType(A) == TypeInfo2.inferScalarType(A));
-    for (VPUser *U : A->users()) {
-      auto *R = cast<VPRecipeBase>(U);
-      for (VPValue *VPV : R->definedValues())
-        assert(TypeInfo.inferScalarType(VPV) == TypeInfo2.inferScalarType(VPV));
-    }
-#endif
   }
 
   // Simplify (X && Y) || (X && !Y) -> X.
@@ -1024,13 +1012,28 @@ void VPlanTransforms::simplifyRecipes(VPlan &Plan, Type &CanonicalIVTy) {
     VPRecipeBase *R = Worklist.pop_back_val();
     if (VPValue *Result = simplifyRecipe(*R, TypeInfo)) {
       R->getVPSingleValue()->replaceAllUsesWith(Result);
+      TypeInfo.erase(R->getVPSingleValue());
       R->eraseFromParent();
       if (VPRecipeBase *ResultR = Result->getDefiningRecipe())
         Worklist.insert(ResultR);
       for (VPUser *U : Result->users())
         if (auto *UR = dyn_cast<VPRecipeBase>(U))
           if (UR != R)
             Worklist.insert(UR);
+
+#ifndef NDEBUG
+      // Verify that the cached type info is for both Result and its users is
+      // still accurate by comparing it to freshly computed types.
+      VPTypeAnalysis TypeInfo2(&CanonicalIVTy);
+      assert(TypeInfo.inferScalarType(Result) ==
+             TypeInfo2.inferScalarType(Result));
+      for (VPUser *U : Result->users()) {
+        auto *R = cast<VPRecipeBase>(U);
+        for (VPValue *VPV : R->definedValues())
+          assert(TypeInfo.inferScalarType(VPV) ==
+                 TypeInfo2.inferScalarType(VPV));
+      }
+#endif
     }
   }
 }