Support mapping Instructions to VPValues in VPRecipeBuilder

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9047,16 +9047,16 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
     for (Instruction &I : drop_end(BB->instructionsWithoutDebug(false))) {
       Instruction *Instr = &I;
 
-      // A special case. Mapping handled in
-      // VPRecipeBuilder::getVPValueOrAddLiveIn().
+      // Special case: Handle extractvalues from struct ret calls.
       if (auto *ExtractValue = dyn_cast<ExtractValueInst>(Instr)) {
-        bool IsUniform = LoopVectorizationPlanner::getDecisionAndClampRange(
-            [&](ElementCount VF) {
-              return CM.isUniformAfterVectorization(ExtractValue, VF);
-            },
-            Range);
-        if (!IsUniform)
+        if (auto *CI = dyn_cast<CallInst>(ExtractValue->getAggregateOperand())) {
+          auto *R = RecipeBuilder.getRecipe(cast<Instruction>(CI));
+          assert(R->getNumDefinedValues() ==
+               cast<StructType>(CI->getType())->getNumElements());
+          unsigned Idx = ExtractValue->getIndices()[0];
+          RecipeBuilder.setRecipe(Instr, R->getVPValue(Idx));
           continue;
+        }
       }
 
       SmallVector<VPValue *, 4> Operands;

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

@@ -54,9 +54,11 @@ class VPRecipeBuilder {
   EdgeMaskCacheTy EdgeMaskCache;
   BlockMaskCacheTy BlockMaskCache;
 
+  using RecipeOrResult = PointerUnion<VPRecipeBase*, VPValue*>;
+
   // VPlan construction support: Hold a mapping from ingredients to
   // their recipe.
-  DenseMap<Instruction *, VPRecipeBase *> Ingredient2Recipe;
+  DenseMap<Instruction *, RecipeOrResult> Ingredient2Recipe;
 
   /// Cross-iteration reduction & first-order recurrence phis for which we need
   /// to add the incoming value from the backedge after all recipes have been
@@ -132,6 +134,14 @@ class VPRecipeBuilder {
     Ingredient2Recipe[I] = R;
   }
 
+  // Set the recipe value for a given ingredient.
+  void setRecipe(Instruction* I, VPValue* RecipeResult) {
+    assert(!Ingredient2Recipe.contains(I) &&
+           "Cannot reset recipe for instruction.");
+    assert(RecipeResult->getDefiningRecipe() && "Value must be defined by a recipe.");
+    Ingredient2Recipe[I] = RecipeResult;
+  }
+
   /// Create the mask for the vector loop header block.
   void createHeaderMask();
 
@@ -158,9 +168,12 @@ class VPRecipeBuilder {
   VPRecipeBase *getRecipe(Instruction *I) {
     assert(Ingredient2Recipe.count(I) &&
            "Recording this ingredients recipe was not requested");
-    assert(Ingredient2Recipe[I] != nullptr &&
+    assert(Ingredient2Recipe[I] &&
            "Ingredient doesn't have a recipe");
-    return Ingredient2Recipe[I];
+    auto RecipeInfo = Ingredient2Recipe[I];
+    if (auto* R = RecipeInfo.dyn_cast<VPRecipeBase*>())
+      return R;
+    return RecipeInfo.get<VPValue*>()->getDefiningRecipe();
   }
 
   /// Build a VPReplicationRecipe for \p I. If it is predicated, add the mask as
@@ -179,17 +192,10 @@ class VPRecipeBuilder {
 
   VPValue *getVPValueOrAddLiveIn(Value *V) {
     if (auto *I = dyn_cast<Instruction>(V)) {
-      if (auto *R = Ingredient2Recipe.lookup(I))
-        return R->getVPSingleValue();
-    }
-    // Ugh: Not sure where to handle this :(
-    if (auto *EVI = dyn_cast<ExtractValueInst>(V)) {
-      Value *AggOp = EVI->getAggregateOperand();
-      if (auto *R = getRecipe(cast<Instruction>(AggOp))) {
-        assert(R->getNumDefinedValues() ==
-               cast<StructType>(AggOp->getType())->getNumElements());
-        unsigned Idx = EVI->getIndices()[0];
-        return R->getVPValue(Idx);
+      if (auto RecipeInfo = Ingredient2Recipe.lookup(I)) {
+        if (auto* R = RecipeInfo.dyn_cast<VPRecipeBase*>())
+          return R->getVPSingleValue();
+        return RecipeInfo.get<VPValue*>();
       }
     }
     return Plan.getOrAddLiveIn(V);