[llvm] [RISCV] Improve stack usage reporting accuracy in the cases when the call frame is not reserved at function entries. (PR #189990)

[via llvm-commits]

https://github.com/zinechant created https://github.com/llvm/llvm-project/pull/189990

Refactors stack size calculation in AsmPrinter by introducing a new virtual method, calculateStackSize. This allows for target-specific overrides.

The RISC-V target overrides this method to conditionally add MaxCallFrameSize to the total stack size when a reserved call frame is not used. This ensures a more accurate stack usage report for --fstack-usage and the .stack_sizes_section.

fixes #189085

>From 30bf7ad78eeb0beda609fe42e1cab5663b89752d Mon Sep 17 00:00:00 2001
From: Chen Zou <chenzou at google.com>
Date: Wed, 1 Apr 2026 01:18:27 +0000
Subject: [PATCH] [RISCV] Refactors stack size calculation in AsmPrinter by
 introducing a new virtual method, calculateStackSize. This allows for
 target-specific overrides.

    The RISC-V target overrides this method to conditionally add MaxCallFrameSize to the total stack size when a reserved call frame is not used. This e
nsures a more accurate stack usage report for --fstack-usage and the .stack_sizes_section.
---
 llvm/include/llvm/CodeGen/AsmPrinter.h     |  6 ++++++
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp | 13 ++++++++-----
 llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp  | 12 ++++++++++++
 llvm/test/CodeGen/RISCV/alloca.ll          |  6 ++++++
 4 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index b160de2f060d6..73d60e7d9621b 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -476,6 +476,12 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
 
   void emitFrameAlloc(const MachineInstr &MI);
 
+  // Helper function to allow potentially post-fix the stack size of a function
+  // via target-specific override.
+  // Will be called by emitStackUsage or emitStackSizeSection, which usually
+  // wouldn't require target-specific overrides.
+  virtual uint64_t calculateStackSize(const MachineFunction &MF);
+
   void emitStackSizeSection(const MachineFunction &MF);
 
   void emitStackUsage(const MachineFunction &MF);
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 5fdf88831892f..2ff2583af3f5a 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1702,6 +1702,11 @@ void AsmPrinter::emitPseudoProbe(const MachineInstr &MI) {
   }
 }
 
+uint64_t AsmPrinter::calculateStackSize(const MachineFunction &MF) {
+  const MachineFrameInfo &FrameInfo = MF.getFrameInfo();
+  return FrameInfo.getStackSize() + FrameInfo.getUnsafeStackSize();
+}
+
 void AsmPrinter::emitStackSizeSection(const MachineFunction &MF) {
   if (!MF.getTarget().Options.EmitStackSizeSection)
     return;
@@ -1720,8 +1725,7 @@ void AsmPrinter::emitStackSizeSection(const MachineFunction &MF) {
   OutStreamer->switchSection(StackSizeSection);
 
   const MCSymbol *FunctionSymbol = getFunctionBegin();
-  uint64_t StackSize =
-      FrameInfo.getStackSize() + FrameInfo.getUnsafeStackSize();
+  const uint64_t StackSize = calculateStackSize(MF);
   OutStreamer->emitSymbolValue(FunctionSymbol, TM.getProgramPointerSize());
   OutStreamer->emitULEB128IntValue(StackSize);
 
@@ -1737,9 +1741,6 @@ void AsmPrinter::emitStackUsage(const MachineFunction &MF) {
   if (OutputFilename.empty())
     return;
 
-  const MachineFrameInfo &FrameInfo = MF.getFrameInfo();
-  uint64_t StackSize =
-      FrameInfo.getStackSize() + FrameInfo.getUnsafeStackSize();
 
   if (StackUsageStream == nullptr) {
     std::error_code EC;
@@ -1756,6 +1757,8 @@ void AsmPrinter::emitStackUsage(const MachineFunction &MF) {
   else
     *StackUsageStream << MF.getFunction().getParent()->getName();
 
+  const uint64_t StackSize = calculateStackSize(MF);
+  const MachineFrameInfo &FrameInfo = MF.getFrameInfo();
   *StackUsageStream << ':' << MF.getName() << '\t' << StackSize << '\t';
   if (FrameInfo.hasVarSizedObjects())
     *StackUsageStream << "dynamic\n";
diff --git a/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp b/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
index eb15227a72a83..e7af7bf01bc28 100644
--- a/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
+++ b/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
@@ -43,6 +43,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/RISCVISAInfo.h"
 #include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
+#include <cstdint>
 
 using namespace llvm;
 
@@ -83,6 +84,8 @@ class RISCVAsmPrinter : public AsmPrinter {
 
   void emitInstruction(const MachineInstr *MI) override;
 
+  uint64_t calculateStackSize(const MachineFunction &MF) override;
+
   void emitMachineConstantPoolValue(MachineConstantPoolValue *MCPV) override;
 
   bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
@@ -358,6 +361,15 @@ void RISCVAsmPrinter::emitInstruction(const MachineInstr *MI) {
   EmitToStreamer(*OutStreamer, OutInst);
 }
 
+uint64_t RISCVAsmPrinter::calculateStackSize(const MachineFunction &MF) {
+  uint64_t StackSize = AsmPrinter::calculateStackSize(MF);
+  const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
+  if (!TFI->hasReservedCallFrame(MF)) {
+    StackSize += MF.getFrameInfo().getMaxCallFrameSize();
+  }
+  return StackSize;
+}
+
 bool RISCVAsmPrinter::PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
                                       const char *ExtraCode, raw_ostream &OS) {
   // First try the generic code, which knows about modifiers like 'c' and 'n'.
diff --git a/llvm/test/CodeGen/RISCV/alloca.ll b/llvm/test/CodeGen/RISCV/alloca.ll
index 9ea5471e4c633..f8315f27038ae 100644
--- a/llvm/test/CodeGen/RISCV/alloca.ll
+++ b/llvm/test/CodeGen/RISCV/alloca.ll
@@ -1,6 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV32I
+; RUN: llc -mtriple=riscv32 -verify-machineinstrs -stack-usage-file=%t.su < %s
+; RUN: FileCheck --input-file=%t.su %s --check-prefix=SU
 
 declare void @notdead(ptr)
 
@@ -100,3 +102,7 @@ define void @alloca_callframe(i32 %n) nounwind {
                   i32 9, i32 10, i32 11, i32 12)
   ret void
 }
+
+; SU: simple_alloca{{.*}}16{{.*}}dynamic
+; SU: scoped_alloca{{.*}}16{{.*}}dynamic
+; SU: alloca_callframe{{.*}}32{{.*}}dynamic