From 1df9e6cdae61f8c076662365a8f97f2dd366a88e Mon Sep 17 00:00:00 2001
From: zhoujing <jing.zhou@terapines.com>
Date: Tue, 25 Jul 2023 10:22:33 +0800
Subject: [PATCH] [VENTUS][RISCV][feat] Initialize adding ventus resource usage
 information to ELF file

---
 llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp   | 26 +++++++++++++++++++++
 llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp |  5 ++--
 llvm/lib/Target/RISCV/RISCVRegisterInfo.h   |  2 +-
 llvm/lib/Target/RISCV/VentusProgramInfo.h   |  6 ++---
 4 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp b/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
index 879d7fc71aa6..37f74bc0373d 100644
--- a/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
+++ b/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
@@ -15,8 +15,10 @@
 #include "MCTargetDesc/RISCVMCExpr.h"
 #include "MCTargetDesc/RISCVTargetStreamer.h"
 #include "RISCV.h"
+#include "RISCVMachineFunctionInfo.h"
 #include "RISCVTargetMachine.h"
 #include "TargetInfo/RISCVTargetInfo.h"
+#include "VentusProgramInfo.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/CodeGen/AsmPrinter.h"
@@ -47,6 +49,7 @@ namespace {
 class RISCVAsmPrinter : public AsmPrinter {
   const MCSubtargetInfo *MCSTI;
   const RISCVSubtarget *STI;
+  VentusProgramInfo CurrentProgramInfo;
 
 public:
   explicit RISCVAsmPrinter(TargetMachine &TM,
@@ -65,6 +68,7 @@ public:
                              const char *ExtraCode, raw_ostream &OS) override;
 
   void EmitToStreamer(MCStreamer &S, const MCInst &Inst);
+  void getVentusProgramInfo(VentusProgramInfo &Out, const MachineFunction &MF);
   bool emitPseudoExpansionLowering(MCStreamer &OutStreamer,
                                    const MachineInstr *MI);
 
@@ -194,6 +198,17 @@ bool RISCVAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   NewSTI.setFeatureBits(MF.getSubtarget().getFeatureBits());
   MCSTI = &NewSTI;
   STI = &MF.getSubtarget<RISCVSubtarget>();
+  CurrentProgramInfo = VentusProgramInfo();
+  if (MF.getInfo<RISCVMachineFunctionInfo>()->isEntryFunction()) {
+    getVentusProgramInfo(CurrentProgramInfo, MF);
+    MCSectionELF *ResourceSection = OutContext.getELFSection(
+        ".rodata.ventus.resource", ELF::SHT_PROGBITS, ELF::SHF_WRITE);
+    OutStreamer->switchSection(ResourceSection);
+    OutStreamer->emitInt16(CurrentProgramInfo.VGPRUsage);
+    OutStreamer->emitInt16(CurrentProgramInfo.SGPRUsage);
+    OutStreamer->emitInt16(CurrentProgramInfo.LDSMemory);
+    OutStreamer->emitInt16(CurrentProgramInfo.PDSMemory);
+  }
 
   SetupMachineFunction(MF);
   emitFunctionBody();
@@ -210,6 +225,17 @@ void RISCVAsmPrinter::emitStartOfAsmFile(Module &M) {
     emitAttributes();
 }
 
+void RISCVAsmPrinter::getVentusProgramInfo(VentusProgramInfo &Out,
+                                           const MachineFunction &MF) {
+  const RISCVSubtarget &ST = MF.getSubtarget<RISCVSubtarget>();
+  const RISCVRegisterInfo *RI = ST.getRegisterInfo();
+  Out.SGPRUsage =
+      RI->getUsedRegistersNum(MF.getRegInfo(), &RISCV::VGPRRegClass, MF);
+  Out.VGPRUsage =
+      RI->getUsedRegistersNum(MF.getRegInfo(), &RISCV::GPRRegClass, MF);
+  // TODO:: Add LDS/PDS calculation
+}
+
 void RISCVAsmPrinter::emitEndOfAsmFile(Module &M) {
   RISCVTargetStreamer &RTS =
       static_cast<RISCVTargetStreamer &>(*OutStreamer->getTargetStreamer());
diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
index 43db6b5f4dd4..7491e2220412 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
@@ -170,7 +170,7 @@ MCRegister RISCVRegisterInfo::findUnusedRegister(const MachineRegisterInfo &MRI,
   return MCRegister();
 }
 
-uint32_t RISCVRegisterInfo::getUnusedRegistersNum(const MachineRegisterInfo &MRI,
+uint32_t RISCVRegisterInfo::getUsedRegistersNum(const MachineRegisterInfo &MRI,
                               const TargetRegisterClass *RC,
                               const MachineFunction &MF) const {
   auto TotalRegNum = std::distance(RC->begin(), RC->end());
@@ -178,7 +178,8 @@ uint32_t RISCVRegisterInfo::getUnusedRegistersNum(const MachineRegisterInfo &MRI
   for (MCRegister Reg : *RC)
     if (MRI.isPhysRegUsed(Reg))
       UsedRegNum++;
-  return TotalRegNum - UsedRegNum;
+  assert(UsedRegNum <= TotalRegNum && "Register using overflow!");
+  return UsedRegNum;
 }
 
 bool RISCVRegisterInfo::isSGPRReg(const MachineRegisterInfo &MRI,
diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.h b/llvm/lib/Target/RISCV/RISCVRegisterInfo.h
index 80ebdf9224f1..f0dd8a79714b 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.h
+++ b/llvm/lib/Target/RISCV/RISCVRegisterInfo.h
@@ -123,7 +123,7 @@ struct RISCVRegisterInfo : public RISCVGenRegisterInfo {
                                 const MachineFunction &MF,
                                 bool ReserveHighestVGPR = false) const;
   
-  uint32_t getUnusedRegistersNum(const MachineRegisterInfo &MRI,
+  uint32_t getUsedRegistersNum(const MachineRegisterInfo &MRI,
                                 const TargetRegisterClass *RC,
                                 const MachineFunction &MF) const;
                                 
diff --git a/llvm/lib/Target/RISCV/VentusProgramInfo.h b/llvm/lib/Target/RISCV/VentusProgramInfo.h
index e7bbaa8739a2..93c0799205e4 100644
--- a/llvm/lib/Target/RISCV/VentusProgramInfo.h
+++ b/llvm/lib/Target/RISCV/VentusProgramInfo.h
@@ -19,9 +19,9 @@ namespace llvm {
 
   struct VentusProgramInfo {
 
-    uint32_t VGPRUsage = 265; //The number of VGPRS which has been used
-    uint32_t SGPRUsage = 265; //The number of SGPRS which has been used
-    uint32_t LDSMemory = 0; //The number of VGPRS which has been used
+    uint32_t VGPRUsage = 256; //The number of VGPRS 
+    uint32_t SGPRUsage = 64; //The number of SGPRS which has been used
+    uint32_t LDSMemory = 1 << 12; //The number of VGPRS which has been used
     uint32_t PDSMemory = 1 << 10; //The number of VGPRS which has been used
 
     VentusProgramInfo() = default;