[llvm] [SLSR] Skipping rewriting based on liveness (PR #218470)

[Igor Wodiany via llvm-commits]

================
@@ -1411,6 +1423,380 @@ bool StraightLineStrengthReduceLegacyPass::runOnFunction(Function &F) {
   return StraightLineStrengthReduce(DL, DT, SE, TTI).runOnFunction(F);
 }
 
+namespace {
+
+class RewriteFilter {
+  // Recondiser rewriting a basic block holding many distinct SLSR bases.
+
+  // Each basis contributes one extended live range no matter how many
+  // candidates are rewritten against it in a basic block.
+  // If the number of bases is above this threshold do a check if the
+  // liveness of the block has increase a lot by SLSR.
+  static constexpr unsigned MinDistinctBasisesToFilter = 16;
+
+public:
+  using Candidate = StraightLineStrengthReduce::Candidate;
+
+  RewriteFilter(const Function *F,
+                DenseMap<Instruction *, Candidate *> &PickedCandidateMap,
+                const TargetTransformInfo *TTI)
+      : F(F), PickedCandidateMap(PickedCandidateMap), TTI(TTI) {}
+
+  // Return the candidates whose rewrite would push their block's register
+  // pressure past what the target can allocate.
+  DenseSet<const Instruction *> run() {
+    DenseSet<const Instruction *> InstsToSkip;
+    if (!EnableRewriteFilter || PickedCandidateMap.empty())
+      return InstsToSkip;
+
+    // Without a budget there is nothing to compare the pressure against, so
+    // check for one before paying for the liveness and pressure analyses.
+    std::optional<unsigned> Budget = getRegisterBudget();
+    if (!Budget)
+      return InstsToSkip;
+
+    buildBBToNumCandsAndBasises(PickedCandidateMap);
+    if (MaxNumBasisesInBB <= MinDistinctBasisesToFilter)
+      return InstsToSkip;
+
+    // Compute live-in and live-out of each BB in CFG
+    buildBBToLiveness(*F);
+
+    DEBUG_SLSR_REWRITE_FILTER(dbgs() << "-- Max liveness of BBs -- \n");
+    SmallPtrSet<const BasicBlock *, 8> BBsToSkip;
+    for (auto &BB : *F) {
+
+      auto It = BBToNumCandsAndBasises.find(&BB);
+      if (It == BBToNumCandsAndBasises.end() ||
+          It->second.second <= MinDistinctBasisesToFilter)
+        continue;
+
+      const BlockLiveness &BL = getLiveness(&BB);
+      auto [MaxLiveness, MaxLivenessWithSLSR] =
+          maxLivenessInBlockBackward(BB, BL.LiveIn, BL.LiveOut);
+      DEBUG_SLSR_REWRITE_FILTER(dbgs() << "MaxLiveness:" << BB.getName()
+                                       << ": (" << MaxLiveness << ", "
+                                       << MaxLivenessWithSLSR << ")" << "\n");
+
+      if (!rewriteWouldOverflowBudget(MaxLiveness, MaxLivenessWithSLSR,
+                                      *Budget))
+        continue;
+
+      DEBUG_SLSR_REWRITE_FILTER(
+          dbgs() << "Skipping BB from SLSR: " << BB.getName() << "\n");
+      BBsToSkip.insert(&BB);
+    } // Done with BBs
+
+    // One pass over the candidates rather than over the instructions of every
+    // skipped block. A skipped block is one of the larger blocks in the
+    // function, while the candidate map is small by comparison.
+    for (const auto &It : PickedCandidateMap)
+      if (BBsToSkip.contains(It.first->getParent()) &&
+          InstsToSkip.insert(It.first).second)
+        NumFilteredCandidates++;
+
+    return InstsToSkip;
+  }
+
+private:
+  const Function *F;
+  DenseMap<Instruction *, Candidate *> &PickedCandidateMap;
+  const TargetTransformInfo *TTI;
+
+  DenseMap<const BasicBlock *, std::pair<unsigned, unsigned>>
+      BBToNumCandsAndBasises;
+  unsigned MaxNumBasisesInBB = 0;
+
+  using ValueSet = SmallPtrSet<const Value *, 32>;
+  struct BlockLiveness {
+    ValueSet LiveIn;
+    ValueSet LiveOut;
+  };
+
+  DenseMap<const BasicBlock *, BlockLiveness> BBToLiveness;
+  // The liveness scan asks for the weight of every live value at every
+  // instruction, so memoize on the type, which is all the weight depends on.
+  mutable DenseMap<Type *, unsigned> WeightCache;
+
+  // Liveness is only computed for functions that pass the candidate-count
+  // gate. Blocks of the remaining functions read as having nothing live across
+  // their boundaries, which is why no rewrite is suppressed in that case.
+  const BlockLiveness &getLiveness(const BasicBlock *BB) const {
+    static const BlockLiveness Empty;
+    auto It = BBToLiveness.find(BB);
+    return It == BBToLiveness.end() ? Empty : It->second;
+  }
+
+  std::optional<unsigned> getRegisterBudget() const {
+    std::optional<unsigned> Budget = TTI->getRegisterBudget(*F);
+    if (Budget && *Budget == 0)
+      return std::nullopt;
+    return Budget;
+  }
+
+  // Return true if rewriting every candidate in a block, taking its peak
+  // pressure from \p Before to \p After, would ask the allocator for more
+  // registers than \p Budget.
+  bool rewriteWouldOverflowBudget(unsigned Before, unsigned After,
+                                  unsigned Budget) const {
+    // Leave the allocator some slack: it also has to satisfy register class
+    // and ABI constraints that this estimate knows nothing about.
+    constexpr double SafeRatio = 0.9;
+    unsigned SafeBudget = static_cast<unsigned>(Budget * SafeRatio);
+
+    // There is headroom, so however much the rewrite adds is irrelevant.
+    if (After <= SafeBudget)
+      return false;
+
+    // The rewrite is what takes the block over.
+    if (Before <= SafeBudget)
+      return true;
+
+    // Already over budget. SLSR can still lower pressure here, so only refuse
+    // rewrites that make it meaningfully worse.
+    constexpr unsigned AbsDelta = 4;
----------------
IgWod wrote:

Isn't it potentially too specific to AMDGPU?

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