[LoopVectorize][NFC] Simplify ScaledReductionExitInstrs map

[david-arm]

For the following variable

DenseMap<const Instruction *, std::pair<PartialReductionChain, unsigned>>
ScaledReductionExitInstrs;

we never actually need the PartialReductionChain when using the map. I've cleaned this up so that this now becomes

DenseMap<const Instruction *, unsigned> ScaledReductionMap;

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: David Sherwood (david-arm)

Changes

For the following variable

DenseMap<const Instruction *, std::pair<PartialReductionChain, unsigned>>
ScaledReductionExitInstrs;

we never actually need the PartialReductionChain when using the map. I've cleaned this up so that this now becomes

DenseMap<const Instruction *, unsigned> ScaledReductionExitInstrs;

---

Full diff: https://github.com/llvm/llvm-project/pull/123368.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+4-3)
• (modified) llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h (+2-3)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 6df11abda9e988..a8da280438b998 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8821,7 +8821,8 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
     PartialReductionChain Chain = Pair.first;
     if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
         ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB))
-      ScaledReductionExitInstrs.insert(std::make_pair(Chain.Reduction, Pair));
+      ScaledReductionExitInstrs.insert(
+          std::make_pair(Chain.Reduction, Pair.second));
   }
 }
 
@@ -8913,9 +8914,9 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
              Phi->getIncomingValueForBlock(OrigLoop->getLoopPreheader()));
 
       // If the PHI is used by a partial reduction, set the scale factor.
-      std::optional<std::pair<PartialReductionChain, unsigned>> Pair =
+      std::optional<unsigned> Scale =
           getScaledReductionForInstr(RdxDesc.getLoopExitInstr());
-      unsigned ScaleFactor = Pair ? Pair->second : 1;
+      unsigned ScaleFactor = Scale ? *Scale : 1;
       PhiRecipe = new VPReductionPHIRecipe(
           Phi, RdxDesc, *StartV, CM.isInLoopReduction(Phi),
           CM.useOrderedReductions(RdxDesc), ScaleFactor);
diff --git a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
index cf653e2d3e6584..492a3b8dfad9f3 100644
--- a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
+++ b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
@@ -88,8 +88,7 @@ class VPRecipeBuilder {
 
   /// The set of reduction exit instructions that will be scaled to
   /// a smaller VF via partial reductions, paired with the scaling factor.
-  DenseMap<const Instruction *, std::pair<PartialReductionChain, unsigned>>
-      ScaledReductionExitInstrs;
+  DenseMap<const Instruction *, unsigned> ScaledReductionExitInstrs;
 
   /// Check if \p I can be widened at the start of \p Range and possibly
   /// decrease the range such that the returned value holds for the entire \p
@@ -157,7 +156,7 @@ class VPRecipeBuilder {
       : Plan(Plan), OrigLoop(OrigLoop), TLI(TLI), TTI(TTI), Legal(Legal),
         CM(CM), PSE(PSE), Builder(Builder) {}
 
-  std::optional<std::pair<PartialReductionChain, unsigned>>
+  std::optional<unsigned>
   getScaledReductionForInstr(const Instruction *ExitInst) {
     auto It = ScaledReductionExitInstrs.find(ExitInst);
     return It == ScaledReductionExitInstrs.end()

[fhahn]

Should the name and comment also be updated now, this now just maps scaled reduction exit instructions to their scaling factors?

[david-arm]

I'm not sure what name I would change it to? Something like ScaledReductionMap? ScalingFactorMap?

[david-arm]

Done. I've tried my best. :)

[SamTebbs33]

LGTM with Florian's suggestions.

[david-arm]

Rebase + address review comments

[fhahn]

LGTM, thanks!

[fhahn]

nit: could also be folded into a single line setting unsigned ScaleFactor directly.

[david-arm]

Good suggestion. Done!