[llvm] AMDGPU gfx12: Add _dvgpr$ symbols for dynamic VGPRs (PR #148251)

[Tim Renouf via llvm-commits]

https://github.com/trenouf created https://github.com/llvm/llvm-project/pull/148251

For each function with the AMDGPU_CS_Chain calling convention, with dynamic VGPRs enabled, add a _dvgpr$ symbol, with the value of the function symbol, plus an offset encoding one less than the number of VGPR blocks used by the function (16 VGPRs per block, no more than 128) in bits 5..3 of the symbol value. This is used by a front-end to have functions that are chained rather than called, and a dispatcher that dynamically resizes the VGPR count before dispatching to a function.

>From 8d0fb2e3dcadf754c9d3a3a69716da7a48660783 Mon Sep 17 00:00:00 2001
From: Tim Renouf <tim.renouf at amd.com>
Date: Tue, 7 Jan 2025 20:15:29 +0000
Subject: [PATCH] AMDGPU gfx12: Add _dvgpr$ symbols for dynamic VGPRs

For each function with the AMDGPU_CS_Chain calling convention, with
dynamic VGPRs enabled, add a _dvgpr$ symbol, with the value of the
function symbol, plus an offset encoding one less than the number of
VGPR blocks used by the function (16 VGPRs per block, no more than 128)
in bits 5..3 of the symbol value. This is used by a front-end to have
functions that are chained rather than called, and a dispatcher that
dynamically resizes the VGPR count before dispatching to a function.
---
 llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp | 26 +++++++++++++++++++++
 llvm/test/CodeGen/AMDGPU/dvgpr_sym.ll       | 12 ++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 llvm/test/CodeGen/AMDGPU/dvgpr_sym.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
index 749b9efc81378..00ed5f57967ce 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
@@ -194,6 +194,32 @@ void AMDGPUAsmPrinter::emitFunctionBodyStart() {
     return;
   }
 
+  if (STM.isDynamicVGPREnabled() &&
+      MF->getFunction().getCallingConv() == CallingConv::AMDGPU_CS_Chain) {
+    // Add a _dvgpr$ symbol, with the value of the function symbol, plus an
+    // offset encoding one less than the number of VGPR blocks used by the
+    // function (16 VGPRs per block, no more than 128) in bits 5..3 of the
+    // symbol value. This is used by a front-end to have functions that are
+    // chained rather than called, and a dispatcher that dynamically resizes
+    // the VGPR count before dispatching to a function.
+    ResourceUsage = &getAnalysis<AMDGPUResourceUsageAnalysis>();
+    const AMDGPUResourceUsageAnalysis::SIFunctionResourceInfo &Info =
+        ResourceUsage->getResourceInfo();
+    MCContext &Ctx = MF->getContext();
+    unsigned EncodedNumVGPRs = (Info.NumVGPR - 1) >> 1 & 0x38;
+    MCSymbol *CurPCSym = Ctx.createTempSymbol();
+    OutStreamer->emitLabel(CurPCSym);
+    const MCExpr *DVgprFuncVal = MCBinaryExpr::createAdd(
+        MCSymbolRefExpr::create(CurPCSym, MCSymbolRefExpr::VK_None, Ctx),
+        MCConstantExpr::create(EncodedNumVGPRs, Ctx), Ctx);
+    MCSymbol *DVgprFuncSym =
+        Ctx.getOrCreateSymbol(Twine("_dvgpr$") + MF->getFunction().getName());
+    OutStreamer->emitAssignment(DVgprFuncSym, DVgprFuncVal);
+    cast<MCSymbolELF>(DVgprFuncSym)
+        ->setBinding(
+            cast<MCSymbolELF>(getSymbol(&MF->getFunction()))->getBinding());
+  }
+
   if (!MFI.isEntryFunction())
     return;
 
diff --git a/llvm/test/CodeGen/AMDGPU/dvgpr_sym.ll b/llvm/test/CodeGen/AMDGPU/dvgpr_sym.ll
new file mode 100644
index 0000000000000..992963d304ead
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/dvgpr_sym.ll
@@ -0,0 +1,12 @@
+; Test generation of _dvgpr$ symbol for an amdgpu_cs_chain function with +dynamic-vgpr.
+
+; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1200 -asm-verbose=0 < %s | FileCheck -check-prefixes=DVGPR %s
+
+; DVGPR-LABEL: func:
+; DVGPR: .Ltmp0:
+; DVGPR: .set _dvgpr$func, .Ltmp0+{{[0-9]+}}
+
+define amdgpu_cs_chain void @func() #0 {
+  ret void
+}
+attributes #0 = { "target-features"="+dynamic-vgpr" }