[llvm] [LoopVectorize][NFC] Simplify ScaledReductionExitInstrs map (PR #123368)

[David Sherwood via llvm-commits]

https://github.com/david-arm updated https://github.com/llvm/llvm-project/pull/123368

>From e94fb885d1b8379fd0b0d2c89597ca91ee2fa2b5 Mon Sep 17 00:00:00 2001
From: David Sherwood <david.sherwood at arm.com>
Date: Fri, 17 Jan 2025 16:29:54 +0000
Subject: [PATCH 1/2] [LoopVectorize][NFC] Simplify ScaledReductionExitInstrs
 map

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;
---
 llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 7 ++++---
 llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h | 5 ++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index f6f80db58cf1b0..206668bc869416 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8711,7 +8711,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));
   }
 }
 
@@ -8803,9 +8804,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()

>From 205e6f2e10bf0958b7b94babe67269e120573719 Mon Sep 17 00:00:00 2001
From: David Sherwood <david.sherwood at arm.com>
Date: Mon, 20 Jan 2025 12:05:35 +0000
Subject: [PATCH 2/2] Address review comments

---
 llvm/lib/Transforms/Vectorize/LoopVectorize.cpp |  9 ++++-----
 llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h | 15 ++++++---------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 206668bc869416..84718d2bc4c5e2 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8711,8 +8711,7 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
     PartialReductionChain Chain = Pair.first;
     if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
         ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB))
-      ScaledReductionExitInstrs.insert(
-          std::make_pair(Chain.Reduction, Pair.second));
+      ScaledReductionMap.insert(std::make_pair(Chain.Reduction, Pair.second));
   }
 }
 
@@ -8805,8 +8804,8 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
 
       // If the PHI is used by a partial reduction, set the scale factor.
       std::optional<unsigned> Scale =
-          getScaledReductionForInstr(RdxDesc.getLoopExitInstr());
-      unsigned ScaleFactor = Scale ? *Scale : 1;
+          getScalingForReduction(RdxDesc.getLoopExitInstr());
+      unsigned ScaleFactor = Scale.value_or(1);
       PhiRecipe = new VPReductionPHIRecipe(
           Phi, RdxDesc, *StartV, CM.isInLoopReduction(Phi),
           CM.useOrderedReductions(RdxDesc), ScaleFactor);
@@ -8841,7 +8840,7 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
   if (isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
     return tryToWidenMemory(Instr, Operands, Range);
 
-  if (getScaledReductionForInstr(Instr))
+  if (getScalingForReduction(Instr))
     return tryToCreatePartialReduction(Instr, Operands);
 
   if (!shouldWiden(Instr, Range))
diff --git a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
index 492a3b8dfad9f3..44745bfd46f891 100644
--- a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
+++ b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
@@ -86,9 +86,8 @@ class VPRecipeBuilder {
   /// created.
   SmallVector<VPHeaderPHIRecipe *, 4> PhisToFix;
 
-  /// 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 *, unsigned> ScaledReductionExitInstrs;
+  /// A mapping of partial reduction exit instructions to their scaling factor.
+  DenseMap<const Instruction *, unsigned> ScaledReductionMap;
 
   /// 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
@@ -156,12 +155,10 @@ class VPRecipeBuilder {
       : Plan(Plan), OrigLoop(OrigLoop), TLI(TLI), TTI(TTI), Legal(Legal),
         CM(CM), PSE(PSE), Builder(Builder) {}
 
-  std::optional<unsigned>
-  getScaledReductionForInstr(const Instruction *ExitInst) {
-    auto It = ScaledReductionExitInstrs.find(ExitInst);
-    return It == ScaledReductionExitInstrs.end()
-               ? std::nullopt
-               : std::make_optional(It->second);
+  std::optional<unsigned> getScalingForReduction(const Instruction *ExitInst) {
+    auto It = ScaledReductionMap.find(ExitInst);
+    return It == ScaledReductionMap.end() ? std::nullopt
+                                          : std::make_optional(It->second);
   }
 
   /// Find all possible partial reductions in the loop and track all of those