From eb3ece00c027c8cacbab520a293df034f17bc086 Mon Sep 17 00:00:00 2001
From: Stanislav Mekhanoshin <Stanislav.Mekhanoshin@amd.com>
Date: Thu, 13 Feb 2025 11:03:22 -0800
Subject: [PATCH] [AMDGPU] Move into SIProgramInfo and cache
 getFunctionCodeSize. NFCI.

This moves function as is, improvements to the estimate go into
a subseqent patch.
---
 llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp | 26 +++-----------------
 llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h   |  2 --
 llvm/lib/Target/AMDGPU/SIProgramInfo.cpp    | 27 +++++++++++++++++++++
 llvm/lib/Target/AMDGPU/SIProgramInfo.h      |  6 +++++
 4 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
index 031d8f0560ff2..a8d0bb746d2ef 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
@@ -748,7 +748,7 @@ bool AMDGPUAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
           RI.getSymbol(CurrentFnSym->getName(), RIK::RIK_PrivateSegSize,
                        OutContext, IsLocal)
               ->getVariableValue(),
-          getFunctionCodeSize(MF), MFI);
+          CurrentProgramInfo.getFunctionCodeSize(MF), MFI);
       return false;
     }
 
@@ -757,7 +757,8 @@ bool AMDGPUAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
         CurrentProgramInfo.NumArchVGPR,
         STM.hasMAIInsts() ? CurrentProgramInfo.NumAccVGPR : nullptr,
         CurrentProgramInfo.NumVGPR, CurrentProgramInfo.NumSGPR,
-        CurrentProgramInfo.ScratchSize, getFunctionCodeSize(MF), MFI);
+        CurrentProgramInfo.ScratchSize,
+        CurrentProgramInfo.getFunctionCodeSize(MF), MFI);
 
     OutStreamer->emitRawComment(
       " FloatMode: " + Twine(CurrentProgramInfo.FloatMode), false);
@@ -893,27 +894,6 @@ void AMDGPUAsmPrinter::initializeTargetID(const Module &M) {
   }
 }
 
-uint64_t AMDGPUAsmPrinter::getFunctionCodeSize(const MachineFunction &MF) const {
-  const GCNSubtarget &STM = MF.getSubtarget<GCNSubtarget>();
-  const SIInstrInfo *TII = STM.getInstrInfo();
-
-  uint64_t CodeSize = 0;
-
-  for (const MachineBasicBlock &MBB : MF) {
-    for (const MachineInstr &MI : MBB) {
-      // TODO: CodeSize should account for multiple functions.
-
-      // TODO: Should we count size of debug info?
-      if (MI.isDebugInstr())
-        continue;
-
-      CodeSize += TII->getInstSizeInBytes(MI);
-    }
-  }
-
-  return CodeSize;
-}
-
 // AccumOffset computed for the MCExpr equivalent of:
 // alignTo(std::max(1, NumVGPR), 4) / 4 - 1;
 static const MCExpr *computeAccumOffset(const MCExpr *NumVGPR, MCContext &Ctx) {
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
index cc8c4411805e2..2c959d7dbbd07 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
@@ -50,8 +50,6 @@ class AMDGPUAsmPrinter final : public AsmPrinter {
 
   MCCodeEmitter *DumpCodeInstEmitter = nullptr;
 
-  uint64_t getFunctionCodeSize(const MachineFunction &MF) const;
-
   void getSIProgramInfo(SIProgramInfo &Out, const MachineFunction &MF);
   void getAmdKernelCode(AMDGPU::AMDGPUMCKernelCodeT &Out,
                         const SIProgramInfo &KernelInfo,
diff --git a/llvm/lib/Target/AMDGPU/SIProgramInfo.cpp b/llvm/lib/Target/AMDGPU/SIProgramInfo.cpp
index 212edff097837..5179288084010 100644
--- a/llvm/lib/Target/AMDGPU/SIProgramInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIProgramInfo.cpp
@@ -27,6 +27,8 @@ void SIProgramInfo::reset(const MachineFunction &MF) {
 
   const MCExpr *ZeroExpr = MCConstantExpr::create(0, Ctx);
 
+  CodeSizeInBytes.reset();
+
   VGPRBlocks = ZeroExpr;
   SGPRBlocks = ZeroExpr;
   Priority = 0;
@@ -199,3 +201,28 @@ const MCExpr *SIProgramInfo::getPGMRSrc2(CallingConv::ID CC,
 
   return MCConstantExpr::create(0, Ctx);
 }
+
+uint64_t SIProgramInfo::getFunctionCodeSize(const MachineFunction &MF) {
+  if (CodeSizeInBytes.has_value())
+    return *CodeSizeInBytes;
+
+  const GCNSubtarget &STM = MF.getSubtarget<GCNSubtarget>();
+  const SIInstrInfo *TII = STM.getInstrInfo();
+
+  uint64_t CodeSize = 0;
+
+  for (const MachineBasicBlock &MBB : MF) {
+    for (const MachineInstr &MI : MBB) {
+      // TODO: CodeSize should account for multiple functions.
+
+      // TODO: Should we count size of debug info?
+      if (MI.isDebugInstr())
+        continue;
+
+      CodeSize += TII->getInstSizeInBytes(MI);
+    }
+  }
+
+  CodeSizeInBytes = CodeSize;
+  return CodeSize;
+}
diff --git a/llvm/lib/Target/AMDGPU/SIProgramInfo.h b/llvm/lib/Target/AMDGPU/SIProgramInfo.h
index 37c03d9b637f0..d7087436ae758 100644
--- a/llvm/lib/Target/AMDGPU/SIProgramInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIProgramInfo.h
@@ -19,6 +19,7 @@
 #include "llvm/IR/CallingConv.h"
 #include "llvm/Support/Compiler.h"
 #include <cstdint>
+#include <optional>
 
 namespace llvm {
 
@@ -29,6 +30,8 @@ class MachineFunction;
 
 /// Track resource usage for kernels / entry functions.
 struct LLVM_EXTERNAL_VISIBILITY SIProgramInfo {
+  std::optional<uint64_t> CodeSizeInBytes;
+
   // Fields set in PGM_RSRC1 pm4 packet.
   const MCExpr *VGPRBlocks = nullptr;
   const MCExpr *SGPRBlocks = nullptr;
@@ -97,6 +100,9 @@ struct LLVM_EXTERNAL_VISIBILITY SIProgramInfo {
   // non-MCExpr members.
   void reset(const MachineFunction &MF);
 
+  // Get function code size and cache the value.
+  uint64_t getFunctionCodeSize(const MachineFunction &MF);
+
   /// Compute the value of the ComputePGMRsrc1 register.
   const MCExpr *getComputePGMRSrc1(const GCNSubtarget &ST,
                                    MCContext &Ctx) const;