From 68d07bf34f6b044193c429d9f8cf07500a4775ea Mon Sep 17 00:00:00 2001
From: Wang Pengcheng <wangpengcheng.pp@bytedance.com>
Date: Thu, 7 Mar 2024 14:18:37 +0800
Subject: [PATCH] [RISCV][NFC] Add helpers for RVV register classes

There are two places in tree that use these helpers and there will
be more future usages.

Reviewers: asb, BeMg, lukel97

Reviewed By: BeMg, lukel97

Pull Request: https://github.com/llvm/llvm-project/pull/84144
---
 llvm/lib/Target/RISCV/RISCVRegisterInfo.h    | 22 ++++++++++++++++++++
 llvm/lib/Target/RISCV/RISCVTargetMachine.cpp | 16 +-------------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.h b/llvm/lib/Target/RISCV/RISCVRegisterInfo.h
index e46fe8ecb900..943c4f2627cf 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.h
+++ b/llvm/lib/Target/RISCV/RISCVRegisterInfo.h
@@ -112,11 +112,33 @@ struct RISCVRegisterInfo : public RISCVGenRegisterInfo {
 
   bool doesRegClassHavePseudoInitUndef(
       const TargetRegisterClass *RC) const override {
+    return isVRRegClass(RC);
+  }
+
+  static bool isVRRegClass(const TargetRegisterClass *RC) {
     return RISCV::VRRegClass.hasSubClassEq(RC) ||
            RISCV::VRM2RegClass.hasSubClassEq(RC) ||
            RISCV::VRM4RegClass.hasSubClassEq(RC) ||
            RISCV::VRM8RegClass.hasSubClassEq(RC);
   }
+
+  static bool isVRNRegClass(const TargetRegisterClass *RC) {
+    return 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 bool isRVVRegClass(const TargetRegisterClass *RC) {
+    return isVRRegClass(RC) || isVRNRegClass(RC);
+  }
 };
 }
 
diff --git a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
index 6fe0abaccb9d..ae1a6f179a49 100644
--- a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
+++ b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
@@ -278,21 +278,7 @@ public:
 
 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);
+  return RISCVRegisterInfo::isRVVRegClass(&RC);
 }
 
 static FunctionPass *useDefaultRegisterAllocator() { return nullptr; }