[llvm] [AMDGPU] Optimize PreRARematStage candidate initialization (PR #220924)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Frederik Harwath (frederik-h)

<details>
<summary>Changes</summary>

The computation of the LiveIn/LiveOut sets in
PreRARematStage::ScoredRemat::init showed up as a significant compile time bottleneck in code using kernels from composable_kernel.

Compute the LiveIn/LiveOut sets for all candidates in a single iteration of the all regions and DAG LiveIns/LiveOuts without iterating the candidates.

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


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp (+39-22) 
- (modified) llvm/lib/Target/AMDGPU/GCNSchedStrategy.h (+9-5) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
index 3dd468772e9b1..6c3a893c9c1e2 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp
@@ -1538,7 +1538,9 @@ bool PreRARematStage::initGCNSchedStage() {
   // Collect candidates. We have more restrictions on what we can track here
   // compared to the rematerializer.
   SmallVector<ScoredRemat, 8> Candidates;
-  SmallVector<unsigned> CandidateOrder;
+  DenseMap<Register, unsigned> DefRegToCandIdx;
+  const unsigned NumRegions = DAG.Regions.size();
+
   for (unsigned RegIdx = 0, E = Remater.getNumRegs(); RegIdx < E; ++RegIdx) {
     const Rematerializer::Reg &CandReg = Remater.getReg(RegIdx);
 
@@ -1592,12 +1594,40 @@ bool PreRARematStage::initGCNSchedStage() {
                      }))
       continue;
 
-    MarkedRegs.insert(CandReg.getDefReg());
-    ScoredRemat &Cand = Candidates.emplace_back();
-    Cand.init(RegIdx, FreqInfo, Remater, DAG);
+    Register DefReg = CandReg.getDefReg();
+    MarkedRegs.insert(DefReg);
+    DefRegToCandIdx[DefReg] = Candidates.size();
+    Candidates.emplace_back(RegIdx, NumRegions);
+  }
+
+  // Initialize the LiveIn and LiveOut sets of all
+  // candidates.
+  //
+  // This used to be part of ScoredRemat::init, but this choice
+  // implied iterating all regions for each candidate. The runtime of
+  // the initialization was O(NumCandidates * NumRegions) whereas this
+  // way it is O(NumRegions + NumLiveIns + NumLiveOuts) (assuming
+  // constant time map lookups).
+  for (unsigned I = 0; I < NumRegions; ++I) {
+    for (const auto &[Reg, Mask] : DAG.LiveIns[I]) {
+      if (auto It = DefRegToCandIdx.find(Reg); It != DefRegToCandIdx.end())
+        Candidates[It->second].LiveIn.set(I);
+    }
+    for (const auto &[Reg, Mask] :
+         DAG.RegionLiveOuts.getLiveRegsForRegionIdx(I)) {
+      if (auto It = DefRegToCandIdx.find(Reg); It != DefRegToCandIdx.end())
+        Candidates[It->second].LiveOut.set(I);
+    }
+  }
+
+  // Finish initializing candidates.
+  SmallVector<unsigned> CandidateOrder;
+  for (unsigned CandIdx = 0; CandIdx < Candidates.size(); ++CandIdx) {
+    ScoredRemat &Cand = Candidates[CandIdx];
+    Cand.init(FreqInfo, Remater, DAG);
     Cand.update(TargetRegions, RPTargets, FreqInfo, !TargetOcc);
     if (!Cand.hasNullScore())
-      CandidateOrder.push_back(Candidates.size() - 1);
+      CandidateOrder.push_back(CandIdx);
   }
 
   if (TargetOcc) {
@@ -3052,36 +3082,23 @@ PreRARematStage::ScoredRemat::FreqInfo::FreqInfo(
   }
 }
 
-void PreRARematStage::ScoredRemat::init(RegisterIdx RegIdx,
-                                        const FreqInfo &Freq,
+void PreRARematStage::ScoredRemat::init(const FreqInfo &Freq,
                                         const Rematerializer &Remater,
                                         GCNScheduleDAGMILive &DAG) {
-  this->RegIdx = RegIdx;
-  const unsigned NumRegions = DAG.Regions.size();
-  LiveIn.resize(NumRegions);
-  LiveOut.resize(NumRegions);
-  Live.resize(NumRegions);
-  UnpredictableRPSave.resize(NumRegions);
-
   const Rematerializer::Reg &Reg = Remater.getReg(RegIdx);
   Register DefReg = Reg.getDefReg();
   assert(Reg.Uses.size() == 1 && "expected users in single region");
   const unsigned UseRegion = Reg.Uses.begin()->first;
 
-  // Mark regions in which the rematerializable register is live.
-  for (unsigned I = 0, E = NumRegions; I != E; ++I) {
-    if (DAG.LiveIns[I].contains(DefReg))
-      LiveIn.set(I);
-    if (DAG.RegionLiveOuts.getLiveRegsForRegionIdx(I).contains(DefReg))
-      LiveOut.set(I);
+  Live |= LiveIn;
+  Live |= LiveOut;
 
+  for (unsigned I : Live.set_bits()) {
     // If the register is both unused and live-through in the region, the
     // latter's RP is guaranteed to decrease.
     if (!LiveIn[I] || !LiveOut[I] || I == UseRegion)
       UnpredictableRPSave.set(I);
   }
-  Live |= LiveIn;
-  Live |= LiveOut;
   RPSave.inc(DefReg, LaneBitmask::getNone(), Reg.Mask, DAG.MRI);
 
   // Get frequencies of defining and using regions. A rematerialization from the
diff --git a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
index 2059f4e6479ff..b49f541504965 100644
--- a/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
+++ b/llvm/lib/Target/AMDGPU/GCNSchedStrategy.h
@@ -575,6 +575,10 @@ class PreRARematStage : public GCNSchedStage {
     /// candidate.
     GCNRegPressure RPSave;
 
+    ScoredRemat(RegisterIdx RegIdx, unsigned NumRegions)
+        : RegIdx(RegIdx), LiveIn(NumRegions), LiveOut(NumRegions),
+          Live(NumRegions), UnpredictableRPSave(NumRegions) {}
+
     /// Execution frequency information required by scoring heuristics.
     /// Frequencies are scaled down if they are high to avoid overflow/underflow
     /// when combining them.
@@ -590,11 +594,11 @@ class PreRARematStage : public GCNSchedStage {
       static const uint64_t ScaleFactor = 1024;
     };
 
-    /// Initializes the candidate with state-independent characteristics for
-    /// rematerializable register with index handle \p RegIdx. This doesn't
-    /// update the actual score (call \ref update for this).
-    void init(RegisterIdx RegIdx, const FreqInfo &Freq,
-              const Rematerializer &Remater, GCNScheduleDAGMILive &DAG);
+    /// Initializes the candidate with state-independent characteristics.
+    /// This doesn't update the actual score (call \ref update for this).
+    /// Note: LiveIn/LiveOut must be pre-populated before calling this.
+    void init(const FreqInfo &Freq, const Rematerializer &Remater,
+              GCNScheduleDAGMILive &DAG);
 
     /// Rematerializes the candidate using the \p Remater.
     void rematerialize(Rematerializer &Remater) const;

``````````

</details>


https://github.com/llvm/llvm-project/pull/220924