[llvm] [MISched] Add a hook to override PostRA scheduling policy (PR #115455)

[Pengcheng Wang via llvm-commits]

https://github.com/wangpc-pp updated https://github.com/llvm/llvm-project/pull/115455

>From 481f87ee0f6fed2f2e1dcaa18b5d8097a5aff654 Mon Sep 17 00:00:00 2001
From: Wang Pengcheng <wangpengcheng.pp at bytedance.com>
Date: Fri, 8 Nov 2024 18:12:31 +0800
Subject: [PATCH 1/2] [MISched] Add a hook to override PostRA scheduling policy

PostRA scheduling supports different directions now, but we can
only specify it via command line options.

This patch adds a new hook `overridePostRASchedPolicy` for targets
to override PostRA scheduling policy.

Note that some options like tracking register pressure won't take
effect in PostRA scheduling.
---
 llvm/include/llvm/CodeGen/TargetSubtargetInfo.h | 10 ++++++++++
 llvm/lib/CodeGen/MachineScheduler.cpp           |  6 ++++++
 2 files changed, 16 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/TargetSubtargetInfo.h b/llvm/include/llvm/CodeGen/TargetSubtargetInfo.h
index bfaa6450779ae0..23d86248ff87a7 100644
--- a/llvm/include/llvm/CodeGen/TargetSubtargetInfo.h
+++ b/llvm/include/llvm/CodeGen/TargetSubtargetInfo.h
@@ -232,6 +232,16 @@ class TargetSubtargetInfo : public MCSubtargetInfo {
   virtual void overrideSchedPolicy(MachineSchedPolicy &Policy,
                                    unsigned NumRegionInstrs) const {}
 
+  /// Override generic post-ra scheduling policy within a region.
+  ///
+  /// This is a convenient way for targets that don't provide any custom
+  /// scheduling heuristics (no custom MachineSchedStrategy) to make
+  /// changes to the generic  post-ra scheduling policy.
+  /// Note that some options like tracking register pressure won't take effect
+  /// in post-ra scheduling.
+  virtual void overridePostRASchedPolicy(MachineSchedPolicy &Policy,
+                                         unsigned NumRegionInstrs) const {}
+
   // Perform target-specific adjustments to the latency of a schedule
   // dependency.
   // If a pair of operands is associated with the schedule dependency, DefOpIdx
diff --git a/llvm/lib/CodeGen/MachineScheduler.cpp b/llvm/lib/CodeGen/MachineScheduler.cpp
index 757f49293993e1..89fd883b56f5f0 100644
--- a/llvm/lib/CodeGen/MachineScheduler.cpp
+++ b/llvm/lib/CodeGen/MachineScheduler.cpp
@@ -3897,6 +3897,12 @@ void PostGenericScheduler::initialize(ScheduleDAGMI *Dag) {
 void PostGenericScheduler::initPolicy(MachineBasicBlock::iterator Begin,
                                       MachineBasicBlock::iterator End,
                                       unsigned NumRegionInstrs) {
+  const MachineFunction &MF = *Begin->getMF();
+
+  // Allow the subtarget to override default policy.
+  MF.getSubtarget().overridePostRASchedPolicy(RegionPolicy, NumRegionInstrs);
+
+  // After subtarget overrides, apply command line options.
   if (PostRADirection == MISchedPostRASched::TopDown) {
     RegionPolicy.OnlyTopDown = true;
     RegionPolicy.OnlyBottomUp = false;

>From a6d00c042bb4b38cab0dfa14f1698a791c32a6d0 Mon Sep 17 00:00:00 2001
From: Wang Pengcheng <wangpengcheng.pp at bytedance.com>
Date: Mon, 11 Nov 2024 15:43:29 +0800
Subject: [PATCH 2/2] Add a getNumOccurrences check

---
 llvm/lib/CodeGen/MachineScheduler.cpp | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineScheduler.cpp b/llvm/lib/CodeGen/MachineScheduler.cpp
index 89fd883b56f5f0..d65db91ac1ebd5 100644
--- a/llvm/lib/CodeGen/MachineScheduler.cpp
+++ b/llvm/lib/CodeGen/MachineScheduler.cpp
@@ -3899,19 +3899,26 @@ void PostGenericScheduler::initPolicy(MachineBasicBlock::iterator Begin,
                                       unsigned NumRegionInstrs) {
   const MachineFunction &MF = *Begin->getMF();
 
+  // Default to top-down because it was implemented first and existing targets
+  // expect that behavior by default.
+  RegionPolicy.OnlyTopDown = true;
+  RegionPolicy.OnlyBottomUp = false;
+
   // Allow the subtarget to override default policy.
   MF.getSubtarget().overridePostRASchedPolicy(RegionPolicy, NumRegionInstrs);
 
   // After subtarget overrides, apply command line options.
-  if (PostRADirection == MISchedPostRASched::TopDown) {
-    RegionPolicy.OnlyTopDown = true;
-    RegionPolicy.OnlyBottomUp = false;
-  } else if (PostRADirection == MISchedPostRASched::BottomUp) {
-    RegionPolicy.OnlyTopDown = false;
-    RegionPolicy.OnlyBottomUp = true;
-  } else if (PostRADirection == MISchedPostRASched::Bidirectional) {
-    RegionPolicy.OnlyBottomUp = false;
-    RegionPolicy.OnlyTopDown = false;
+  if (PostRADirection.getNumOccurrences() > 0) {
+    if (PostRADirection == MISchedPostRASched::TopDown) {
+      RegionPolicy.OnlyTopDown = true;
+      RegionPolicy.OnlyBottomUp = false;
+    } else if (PostRADirection == MISchedPostRASched::BottomUp) {
+      RegionPolicy.OnlyTopDown = false;
+      RegionPolicy.OnlyBottomUp = true;
+    } else if (PostRADirection == MISchedPostRASched::Bidirectional) {
+      RegionPolicy.OnlyBottomUp = false;
+      RegionPolicy.OnlyTopDown = false;
+    }
   }
 }