[H2BLB] Hook up our custom scheduling strategy

This commit teaches our TargetPassConfig class how to use our custom node-picking strategy for our pre-register-allocation machine scheduler instance.
PacktPublishing · Dec 30, 2024 · 537bd8f · 537bd8f
1 parent ca6b29a
commit 537bd8f
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/llvm/lib/Target/H2BLB/H2BLBTargetMachine.cpp b/llvm/lib/Target/H2BLB/H2BLBTargetMachine.cpp
@@ -12,23 +12,31 @@
 
 #include "H2BLBTargetMachine.h"
 #include "H2BLB.h"
+#include "H2BLBMachineScheduler.h"
 #include "H2BLBTargetObjectFile.h"
 #include "H2BLBTargetTransformInfo.h"
 #include "TargetInfo/H2BLBTargetInfo.h" // For getTheH2BLBTarget.
 #include "llvm/CodeGen/GlobalISel/IRTranslator.h"
 #include "llvm/CodeGen/GlobalISel/InstructionSelect.h"
 #include "llvm/CodeGen/GlobalISel/Legalizer.h"
 #include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
+#include "llvm/CodeGen/MachineScheduler.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
-#include "llvm/InitializePasses.h" // For initializeGlobalISel.
+#include "llvm/InitializePasses.h"  // For initializeGlobalISel.
 #include "llvm/MC/TargetRegistry.h" // For RegisterTargetMachine.
 #include "llvm/Passes/PassBuilder.h"
-#include "llvm/Support/CodeGen.h"  // For CodeGenOptLevel.
+#include "llvm/Support/CodeGen.h" // For CodeGenOptLevel.
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h" // For LLVM_EXTERNAL_VISIBILITY.
 #include <memory>
 
 using namespace llvm;
 
+static cl::opt<bool>
+    UseCustomSched("h2blb-use-custom-sched", cl::Hidden,
+                   cl::desc("Enable the H2BLB custom scheduler strategy"),
+                   cl::init(true));
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeH2BLBTarget() {
   // Register the target so that external tools can instantiate it.
   RegisterTargetMachine<H2BLBTargetMachine> X(getTheH2BLBTarget());
@@ -153,3 +161,13 @@ void H2BLBPassConfig::addIRPasses() {
   if (getOptLevel() != CodeGenOptLevel::None)
     addPass(createH2BLBSimpleConstantPropagationPassForLegacyPM());
 }
+
+ScheduleDAGInstrs *
+H2BLBPassConfig::createMachineScheduler(MachineSchedContext *C) const {
+  ScheduleDAGMILive *DAG = new ScheduleDAGMILive(
+      C, UseCustomSched ? std::make_unique<H2BLBPreRASchedStrategy>(C)
+                        : std::make_unique<GenericScheduler>(C));
+  // add DAG Mutations here.
+
+  return DAG;
+}
diff --git a/llvm/lib/Target/H2BLB/H2BLBTargetMachine.h b/llvm/lib/Target/H2BLB/H2BLBTargetMachine.h
@@ -61,6 +61,10 @@ class H2BLBPassConfig : public TargetPassConfig {
   void addPreRegBankSelect() override;
   bool addRegBankSelect() override;
   bool addGlobalInstructionSelect() override;
+
+  // Add custom machine scheduler heuristic;
+  ScheduleDAGInstrs *
+  createMachineScheduler(MachineSchedContext *C) const override;
 };
 
 } // end namespace llvm

diff --git a/llvm/test/CodeGen/H2BLB/sched.mir b/llvm/test/CodeGen/H2BLB/sched.mir
@@ -136,9 +136,9 @@ body:             |
     ; GENERIC-NEXT: [[COPY2:%[0-9]+]]:gpr16 = COPY $r3
     ; GENERIC-NEXT: [[LDR16_:%[0-9]+]]:gpr16 = LDR16 [[COPY2]], 0
     ; GENERIC-NEXT: [[ADDi16rr:%[0-9]+]]:gpr16 = ADDi16rr [[COPY]], [[COPY1]]
+    ; GENERIC-NEXT: [[WIDENING_SMUL:%[0-9]+]]:gpr32 = WIDENING_SMUL [[COPY]], [[LDR16_]]
     ; GENERIC-NEXT: [[ADDi16rr1:%[0-9]+]]:gpr16 = ADDi16rr [[ADDi16rr]], [[ADDi16rr]]
     ; GENERIC-NEXT: [[ADDi16rr2:%[0-9]+]]:gpr16 = ADDi16rr [[ADDi16rr1]], [[ADDi16rr1]]
-    ; GENERIC-NEXT: [[WIDENING_SMUL:%[0-9]+]]:gpr32 = WIDENING_SMUL [[COPY]], [[LDR16_]]
     ; GENERIC-NEXT: [[ADDi16rr3:%[0-9]+]]:gpr16 = ADDi16rr [[ADDi16rr2]], [[WIDENING_SMUL]].sub_low16
     ; GENERIC-NEXT: $r1 = COPY [[ADDi16rr3]]
     ; GENERIC-NEXT: RETURN implicit $r0, implicit $r1