From 56a1c307f99ed03a889756add3ecc554e0c4844c Mon Sep 17 00:00:00 2001
From: Piyou Chen <piyou.chen@sifive.com>
Date: Thu, 16 Nov 2023 04:09:08 -0800
Subject: [PATCH] [RISCV] Split regalloc between RVV and other

Enable this flow by -riscv-split-regalloc=1 (default disable), and could designate specific allocator to RVV by
-riscv-rvv-regalloc=<fast|basic|greedy>

It uses the RegClass filter function to decide which regclass need to be processed.

This patch is pre-requirement for supporting PostRA vsetvl insertion pass.
---
 llvm/lib/Target/RISCV/RISCVTargetMachine.cpp | 108 +++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
index 1b2b3fc1dbab4..524c1a5ca50c5 100644
--- a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
+++ b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
@@ -27,6 +27,7 @@
 #include "llvm/CodeGen/MIRParser/MIParser.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/InitializePasses.h"
@@ -89,6 +90,11 @@ static cl::opt<bool>
                            cl::desc("Enable the loop data prefetch pass"),
                            cl::init(true));
 
+static cl::opt<bool>
+    EnableSplitRegAlloc("riscv-split-regalloc", cl::Hidden,
+                        cl::desc("Enable Split RegisterAlloc for RVV"),
+                        cl::init(false));
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeRISCVTarget() {
   RegisterTargetMachine<RISCVTargetMachine> X(getTheRISCV32Target());
   RegisterTargetMachine<RISCVTargetMachine> Y(getTheRISCV64Target());
@@ -253,6 +259,76 @@ bool RISCVTargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
 }
 
 namespace {
+
+class RVVRegisterRegAlloc : public RegisterRegAllocBase<RVVRegisterRegAlloc> {
+public:
+  RVVRegisterRegAlloc(const char *N, const char *D, FunctionPassCtor C)
+      : RegisterRegAllocBase(N, D, C) {}
+};
+
+static bool onlyAllocateRVVReg(const TargetRegisterInfo &TRI,
+                               const TargetRegisterClass &RC) {
+  return RISCV::VRRegClass.hasSubClassEq(&RC) ||
+         RISCV::VRM2RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRM4RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRM8RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN2M1RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN2M2RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN2M4RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN3M1RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN3M2RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN4M1RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN4M2RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN5M1RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN6M1RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN7M1RegClass.hasSubClassEq(&RC) ||
+         RISCV::VRN8M1RegClass.hasSubClassEq(&RC);
+}
+
+static FunctionPass *useDefaultRegisterAllocator() { return nullptr; }
+
+static llvm::once_flag InitializeDefaultRVVRegisterAllocatorFlag;
+
+/// -riscv-rvv-regalloc=<fast|basic|greedy> command line option.
+/// This option could designate the rvv register allocator only.
+/// For example: -riscv-rvv-regalloc=basic
+static cl::opt<RVVRegisterRegAlloc::FunctionPassCtor, false,
+               RegisterPassParser<RVVRegisterRegAlloc>>
+    RVVRegAlloc("riscv-rvv-regalloc", cl::Hidden,
+                cl::init(&useDefaultRegisterAllocator),
+                cl::desc("Register allocator to use for RVV register."));
+
+static void initializeDefaultRVVRegisterAllocatorOnce() {
+  RegisterRegAlloc::FunctionPassCtor Ctor = RVVRegisterRegAlloc::getDefault();
+
+  if (!Ctor) {
+    Ctor = RVVRegAlloc;
+    RVVRegisterRegAlloc::setDefault(RVVRegAlloc);
+  }
+}
+
+static FunctionPass *createBasicRVVRegisterAllocator() {
+  return createBasicRegisterAllocator(onlyAllocateRVVReg);
+}
+
+static FunctionPass *createGreedyRVVRegisterAllocator() {
+  return createGreedyRegisterAllocator(onlyAllocateRVVReg);
+}
+
+static FunctionPass *createFastRVVRegisterAllocator() {
+  return createFastRegisterAllocator(onlyAllocateRVVReg, false);
+}
+
+static RVVRegisterRegAlloc basicRegAllocRVVReg("basic",
+                                               "basic register allocator",
+                                               createBasicRVVRegisterAllocator);
+static RVVRegisterRegAlloc
+    greedyRegAllocRVVReg("greedy", "greedy register allocator",
+                         createGreedyRVVRegisterAllocator);
+
+static RVVRegisterRegAlloc fastRegAllocRVVReg("fast", "fast register allocator",
+                                              createFastRVVRegisterAllocator);
+
 class RISCVPassConfig : public TargetPassConfig {
 public:
   RISCVPassConfig(RISCVTargetMachine &TM, PassManagerBase &PM)
@@ -301,6 +377,9 @@ class RISCVPassConfig : public TargetPassConfig {
   void addPreEmitPass2() override;
   void addPreSched2() override;
   void addMachineSSAOptimization() override;
+  FunctionPass *createRVVRegAllocPass(bool Optimized);
+  bool addRegAssignAndRewriteFast() override;
+  bool addRegAssignAndRewriteOptimized() override;
   void addPreRegAlloc() override;
   void addPostRegAlloc() override;
   void addOptimizedRegAlloc() override;
@@ -312,6 +391,35 @@ TargetPassConfig *RISCVTargetMachine::createPassConfig(PassManagerBase &PM) {
   return new RISCVPassConfig(*this, PM);
 }
 
+FunctionPass *RISCVPassConfig::createRVVRegAllocPass(bool Optimized) {
+  // Initialize the global default.
+  llvm::call_once(InitializeDefaultRVVRegisterAllocatorFlag,
+                  initializeDefaultRVVRegisterAllocatorOnce);
+
+  RegisterRegAlloc::FunctionPassCtor Ctor = RVVRegisterRegAlloc::getDefault();
+  if (Ctor != useDefaultRegisterAllocator)
+    return Ctor();
+
+  if (Optimized)
+    return createGreedyRVVRegisterAllocator();
+
+  return createFastRVVRegisterAllocator();
+}
+
+bool RISCVPassConfig::addRegAssignAndRewriteFast() {
+  if (EnableSplitRegAlloc)
+    addPass(createRVVRegAllocPass(false));
+  return TargetPassConfig::addRegAssignAndRewriteFast();
+}
+
+bool RISCVPassConfig::addRegAssignAndRewriteOptimized() {
+  if (EnableSplitRegAlloc) {
+    addPass(createRVVRegAllocPass(true));
+    addPass(createVirtRegRewriter(false));
+  }
+  return TargetPassConfig::addRegAssignAndRewriteOptimized();
+}
+
 void RISCVPassConfig::addIRPasses() {
   addPass(createAtomicExpandPass());