[AMDGPU] Mark amdgpu_gfx functions as module entry function

[sebastian-ne]

Summary

• Allows lds allocations
• Writes resource usage into COMPUTE_PGM_RSRC1 registers in PAL metadata

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

939–940

One or the other error methods

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

51

Naming is a bit weird for the comment

[sebastian-ne]

Makes sense, thanks

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

51

"the driver" is also unclear

[sebastian-ne]

Change comments to differentiate between calls from hardware and calls from other functions.
I hope this is more clear than “driver”.

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

bool AMDGPUAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   MCContext &Context = getObjFileLowering().getContext();
   // FIXME: This should be an explicit check for Mesa.
   if (!STM.isAmdHsaOS() && !STM.isAmdPalOS()) {
     MCSectionELF *ConfigSection =
         Context.getELFSection(".AMDGPU.config", ELF::SHT_PROGBITS, 0);
     OutStreamer->SwitchSection(ConfigSection);
   }
 
-  if (MFI->isEntryFunction()) {
+  if (MFI->isModuleEntryFunction()) {
     getSIProgramInfo(CurrentProgramInfo, MF);
   } else {
     auto I = CallGraphResourceInfo.insert(
       std::make_pair(&MF.getFunction(), SIFunctionResourceInfo()));
     SIFunctionResourceInfo &Info = I.first->second;
     assert(I.second && "should only be called once per function");
     Info = analyzeResourceUsage(MF);
   }
 
   if (STM.isAmdPalOS()) {
     if (MFI->isEntryFunction())
       EmitPALMetadata(MF, CurrentProgramInfo);
-    else
+    else if (MFI->isModuleEntryFunction())
       emitPALFunctionMetadata(MF);
   } else if (!STM.isAmdHsaOS()) {
     EmitProgramInfoSI(MF, CurrentProgramInfo);
   }
 
   DumpCodeInstEmitter = nullptr;
   if (STM.dumpCode()) {
     // For -dumpcode, get the assembler out of the streamer, even if it does
for (const MachineInstr &MI : MBB) {
       if (MI.isCall()) {
         // Pseudo used just to encode the underlying global. Is there a better
         // way to track this?
 
         const MachineOperand *CalleeOp
           = TII->getNamedOperand(MI, AMDGPU::OpName::callee);
 
         const Function *Callee = getCalleeFunction(*CalleeOp);
-        if (!Callee || Callee->isDeclaration()) {
+        DenseMap<const Function *, SIFunctionResourceInfo>::const_iterator I =
+            CallGraphResourceInfo.end();
+        bool IsExternal = !Callee || Callee->isDeclaration();
+        if (!IsExternal)
+          I = CallGraphResourceInfo.find(Callee);
+
+        if (IsExternal || I == CallGraphResourceInfo.end()) {
+          // Avoid crashing on undefined behavior with an illegal call to a
+          // kernel. If a callsite's calling convention doesn't match the
+          // function's, it's undefined behavior. If the callsite calling
+          // convention does match, that would have errored earlier.
+          // FIXME: The verifier shouldn't allow this.
+          if (!IsExternal &&
+              AMDGPU::isEntryFunctionCC(Callee->getCallingConv()))
+            report_fatal_error("invalid call to entry function");
+
           // If this is a call to an external function, we can't do much. Make
           // conservative guesses.
 
           // 48 SGPRs - vcc, - flat_scr, -xnack
           int MaxSGPRGuess =
             47 - IsaInfo::getNumExtraSGPRs(&ST, true, ST.hasFlatAddressSpace());
           MaxSGPR = std::max(MaxSGPR, MaxSGPRGuess);
           MaxVGPR = std::max(MaxVGPR, 23);
           MaxAGPR = std::max(MaxAGPR, 23);
 
           CalleeFrameSize = std::max(CalleeFrameSize,
             static_cast<uint64_t>(AssumedStackSizeForExternalCall));
 
           Info.UsesVCC = true;
           Info.UsesFlatScratch = ST.hasFlatAddressSpace();
           Info.HasDynamicallySizedStack = true;
         } else {
           // We force CodeGen to run in SCC order, so the callee's register
           // usage etc. should be the cumulative usage of all callees.
 
-          auto I = CallGraphResourceInfo.find(Callee);
-          if (I == CallGraphResourceInfo.end()) {
-            // Avoid crashing on undefined behavior with an illegal call to a
-            // kernel. If a callsite's calling convention doesn't match the
-            // function's, it's undefined behavior. If the callsite calling
-            // convention does match, that would have errored earlier.
-            // FIXME: The verifier shouldn't allow this.
-            if (AMDGPU::isEntryFunctionCC(Callee->getCallingConv()))
-              report_fatal_error("invalid call to entry function");
-
-            llvm_unreachable("callee should have been handled before caller");
-          }
-
           MaxSGPR = std::max(I->second.NumExplicitSGPR - 1, MaxSGPR);
           MaxVGPR = std::max(I->second.NumVGPR - 1, MaxVGPR);
           MaxAGPR = std::max(I->second.NumAGPR - 1, MaxAGPR);
           CalleeFrameSize
             = std::max(I->second.PrivateSegmentSize, CalleeFrameSize);
           Info.UsesVCC |= I->second.UsesVCC;
           Info.UsesFlatScratch |= I->second.UsesFlatScratch;
           Info.HasDynamicallySizedStack |= I->second.HasDynamicallySizedStack;
void AMDGPUAsmPrinter::EmitPALMetadata(const MachineFunction &MF,
   const GCNSubtarget &STM = MF.getSubtarget<GCNSubtarget>();
   if (STM.isWave32())
     MD->setWave32(MF.getFunction().getCallingConv());
 }
 
 void AMDGPUAsmPrinter::emitPALFunctionMetadata(const MachineFunction &MF) {
   auto *MD = getTargetStreamer()->getPALMetadata();
   const MachineFrameInfo &MFI = MF.getFrameInfo();
-  MD->setStackFrameSize(MF, MFI.getStackSize());
+  MD->setFunctionScratchSize(MF, MFI.getStackSize());
+  // Set compute registers
+  MD->setRsrc1(CallingConv::AMDGPU_CS,
+               CurrentProgramInfo.getPGMRSrc1(CallingConv::AMDGPU_CS));
+  MD->setRsrc2(CallingConv::AMDGPU_CS, CurrentProgramInfo.ComputePGMRSrc2);
 }
 
 // This is supposed to be log2(Size)
 static amd_element_byte_size_t getElementByteSizeValue(unsigned Size) {
   switch (Size) {
   case 4:
     return AMD_ELEMENT_4_BYTES;
   case 8:

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

SDValue AMDGPUTargetLowering::LowerGlobalAddress(AMDGPUMachineFunction* MFI,
                                                  SelectionDAG &DAG) const {
 
   const DataLayout &DL = DAG.getDataLayout();
   GlobalAddressSDNode *G = cast<GlobalAddressSDNode>(Op);
   const GlobalValue *GV = G->getGlobal();
 
   if (G->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS ||
       G->getAddressSpace() == AMDGPUAS::REGION_ADDRESS) {
-    if (!MFI->isEntryFunction()) {
+    if (!MFI->isModuleEntryFunction()) {
       SDLoc DL(Op);
       const Function &Fn = DAG.getMachineFunction().getFunction();
       DiagnosticInfoUnsupported BadLDSDecl(
         Fn, "local memory global used by non-kernel function",
         DL.getDebugLoc(), DS_Warning);
       DAG.getContext()->diagnose(BadLDSDecl);
 
       // We currently don't have a way to correctly allocate LDS objects that

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

bool AMDGPULegalizerInfo::legalizeGlobalValue(
   LLT Ty = MRI.getType(DstReg);
   unsigned AS = Ty.getAddressSpace();
 
   const GlobalValue *GV = MI.getOperand(1).getGlobal();
   MachineFunction &MF = B.getMF();
   SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
 
   if (AS == AMDGPUAS::LOCAL_ADDRESS || AS == AMDGPUAS::REGION_ADDRESS) {
-    if (!MFI->isEntryFunction()) {
+    if (!MFI->isModuleEntryFunction()) {
       const Function &Fn = MF.getFunction();
       DiagnosticInfoUnsupported BadLDSDecl(
         Fn, "local memory global used by non-kernel function", MI.getDebugLoc(),
         DS_Warning);
       Fn.getContext().diagnose(BadLDSDecl);
 
       // We currently don't have a way to correctly allocate LDS objects that
       // aren't directly associated with a kernel. We do force inlining of

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

protected:
   /// LDSSize to ensure that dynamic one is aligned accordingly.
   /// The maximal alignment is updated during IR translation or lowering
   /// stages.
   Align DynLDSAlign;
 
   // State of MODE register, assumed FP mode.
   AMDGPU::SIModeRegisterDefaults Mode;
 
-  // Kernels + shaders. i.e. functions called by the driver and not called
+  // Kernels + shaders. i.e. functions called by the hardware and not called
   // by other functions.
   bool IsEntryFunction = false;
 
+  // Entry points called by other functions instead of directly by the hardware.
+  bool IsModuleEntryFunction = false;
+
   bool NoSignedZerosFPMath = false;
 
   // Function may be memory bound.
   bool MemoryBound = false;
 
   // Kernel may need limited waves per EU for better performance.
   bool WaveLimiter = false;
 
public:
   AMDGPU::SIModeRegisterDefaults getMode() const {
     return Mode;
   }
 
   bool isEntryFunction() const {
     return IsEntryFunction;
   }
 
+  bool isModuleEntryFunction() const { return IsModuleEntryFunction; }
+
   bool hasNoSignedZerosFPMath() const {
     return NoSignedZerosFPMath;
   }
 
   bool isMemoryBound() const {
     return MemoryBound;
   }
 

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

 //===-- AMDGPUMachineFunctionInfo.cpp ---------------------------------------=//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPUMachineFunction.h"
 #include "AMDGPUSubtarget.h"
 #include "AMDGPUPerfHintAnalysis.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 
 using namespace llvm;
 
-AMDGPUMachineFunction::AMDGPUMachineFunction(const MachineFunction &MF) :
-  MachineFunctionInfo(),
-  Mode(MF.getFunction()),
-  IsEntryFunction(AMDGPU::isEntryFunctionCC(MF.getFunction().getCallingConv())),
+AMDGPUMachineFunction::AMDGPUMachineFunction(const MachineFunction &MF)
+    : MachineFunctionInfo(), Mode(MF.getFunction()),
+      IsEntryFunction(
+          AMDGPU::isEntryFunctionCC(MF.getFunction().getCallingConv())),
+      IsModuleEntryFunction(
+          AMDGPU::isModuleEntryFunctionCC(MF.getFunction().getCallingConv())),
       NoSignedZerosFPMath(MF.getTarget().Options.NoSignedZerosFPMath) {
   const AMDGPUSubtarget &ST = AMDGPUSubtarget::get(MF);
 
   // FIXME: Should initialize KernArgSize based on ExplicitKernelArgOffset,
   // except reserved size is not correctly aligned.
   const Function &F = MF.getFunction();
 
   Attribute MemBoundAttr = F.getFnAttribute("amdgpu-memory-bound");
   MemoryBound = MemBoundAttr.isStringAttribute() &&

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.h

 bool isGraphics(CallingConv::ID CC);
 
 LLVM_READNONE
 bool isCompute(CallingConv::ID CC);
 
 LLVM_READNONE
 bool isEntryFunctionCC(CallingConv::ID CC);
 
+// These functions are considered entrypoints into the current module, i.e. they
+// are allowed to be called from outside the current module. This is different
+// from isEntryFunctionCC, which is only true for functions that are entered by
+// the hardware. Module entry points include all entry functions but also
+// include functions that can be called from other functions inside or outside
+// the current module. Module entry functions are allowed to allocate LDS.
+LLVM_READNONE
+bool isModuleEntryFunctionCC(CallingConv::ID CC);
+
 // FIXME: Remove this when calling conventions cleaned up
 LLVM_READNONE
 inline bool isKernel(CallingConv::ID CC) {
   switch (CC) {
   case CallingConv::AMDGPU_KERNEL:
   case CallingConv::SPIR_KERNEL:
     return true;
   default:

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

bool isEntryFunctionCC(CallingConv::ID CC) {
   case CallingConv::AMDGPU_HS:
   case CallingConv::AMDGPU_LS:
     return true;
   default:
     return false;
   }
 }
 
+bool isModuleEntryFunctionCC(CallingConv::ID CC) {
+  switch (CC) {
+  case CallingConv::AMDGPU_Gfx:
+    return true;
+  default:
+    return isEntryFunctionCC(CC);
+  }
+}
+
 bool hasXNACK(const MCSubtargetInfo &STI) {
   return STI.getFeatureBits()[AMDGPU::FeatureXNACK];
 }
 
 bool hasSRAMECC(const MCSubtargetInfo &STI) {
   return STI.getFeatureBits()[AMDGPU::FeatureSRAMECC];
 }
 

llvm/lib/Target/AMDGPU/Utils/AMDGPUPALMetadata.h

public:
   // record for logging etc; wave dispatch actually uses the rsrc1 register for
   // the shader stage to determine the number of sgprs to allocate.
   void setNumUsedSgprs(unsigned CC, unsigned Val);
 
   // Set the scratch size in the metadata.
   void setScratchSize(unsigned CC, unsigned Val);
 
   // Set the stack frame size of a function in the metadata.
-  void setStackFrameSize(const MachineFunction &MF, unsigned Val);
+  void setFunctionScratchSize(const MachineFunction &MF, unsigned Val);
 
   // Set the hardware register bit in PAL metadata to enable wave32 on the
   // shader of the given calling convention.
   void setWave32(unsigned CC);
 
   // Emit the accumulated PAL metadata as asm directives.
   // This is called from AMDGPUTargetAsmStreamer::Finish().
   void toString(std::string &S);
private:
   msgpack::MapDocNode getRegisters();
 
   // Reference (create if necessary) the node for the shader functions map.
   msgpack::DocNode &refShaderFunctions();
 
   // Get (create if necessary) the shader functions map.
   msgpack::MapDocNode getShaderFunctions();
 
+  // Get (create if necessary) a function in the shader functions map.
+  msgpack::MapDocNode getShaderFunction(StringRef Name);
+
   // Get (create if necessary) the .hardware_stages entry for the given calling
   // convention.
   msgpack::MapDocNode getHwStage(unsigned CC);
 
   bool setFromLegacyBlob(StringRef Blob);
   bool setFromMsgPackBlob(StringRef Blob);
   void toLegacyBlob(std::string &Blob);
   void toMsgPackBlob(std::string &Blob);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_AMDGPUPALMETADATA_H

llvm/lib/Target/AMDGPU/Utils/AMDGPUPALMetadata.cpp

if (isLegacy()) {
     // Old non-msgpack format.
     setRegister(getScratchSizeKey(CC), Val);
     return;
   }
   // Msgpack format.
   getHwStage(CC)[".scratch_memory_size"] = MsgPackDoc.getNode(Val);
 }
 
-// Set the scratch size in the metadata.
-void AMDGPUPALMetadata::setStackFrameSize(const MachineFunction &MF,
+// Set the stack frame size of a function in the metadata.
+void AMDGPUPALMetadata::setFunctionScratchSize(const MachineFunction &MF,
                                                unsigned Val) {
-  auto Node = MsgPackDoc.getMapNode();
+  auto Node = getShaderFunction(MF.getFunction().getName());
   Node[".stack_frame_size_in_bytes"] = MsgPackDoc.getNode(Val);
-  getShaderFunctions()[MF.getFunction().getName()] = Node;
 }
 
 // Set the hardware register bit in PAL metadata to enable wave32 on the
 // shader of the given calling convention.
 void AMDGPUPALMetadata::setWave32(unsigned CC) {
   switch (CC) {
   case CallingConv::AMDGPU_HS:
     setRegister(PALMD::R_A2D5_VGT_SHADER_STAGES_EN, S_028B54_HS_W32_EN(1));

 
 // Get (create if necessary) the shader functions map.
 msgpack::MapDocNode AMDGPUPALMetadata::getShaderFunctions() {
   if (ShaderFunctions.isEmpty())
     ShaderFunctions = refShaderFunctions();
   return ShaderFunctions.getMap();
 }
 
+// Get (create if necessary) a function in the shader functions map.
+msgpack::MapDocNode AMDGPUPALMetadata::getShaderFunction(StringRef Name) {
+  auto Functions = getShaderFunctions();
+  return Functions[Name].getMap(/*Convert=*/true);
+}
+
 // Return the PAL metadata hardware shader stage name.
 static const char *getStageName(CallingConv::ID CC) {
   switch (CC) {
   case CallingConv::AMDGPU_PS:
     return ".ps";
   case CallingConv::AMDGPU_VS:
     return ".vs";
   case CallingConv::AMDGPU_GS:

llvm/test/CodeGen/AMDGPU/amdpal-callable.ll

 ; RUN: llc -mtriple=amdgcn--amdpal -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SDAG -enable-var-scope %s
-; RUN: llc -mtriple=amdgcn--amdpal -mcpu=tonga -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SDAG -enable-var-scope %s
 ; RUN: llc -mtriple=amdgcn--amdpal -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,SDAG -enable-var-scope %s
 ; RUN: llc -global-isel -mtriple=amdgcn--amdpal -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,GISEL -enable-var-scope %s
 
 declare float @extern_func(float) #0
 declare float @extern_func_many_args(<64 x float>) #0
 
 @funcptr = external hidden unnamed_addr addrspace(4) constant void()*, align 4
 
define amdgpu_gfx float @simple_stack_recurse(float %arg0) #0 {
   %stack = alloca float, i32 4, align 4, addrspace(5)
   store volatile float 2.0, float addrspace(5)* %stack
   %val = load volatile float, float addrspace(5)* %stack
   %res = call amdgpu_gfx float @simple_stack_recurse(float %arg0)
   %add = fadd float %res, %val
   ret float %add
 }
 
+@lds = internal addrspace(3) global [64 x float] undef
+
+define amdgpu_gfx float @simple_lds(float %arg0) #0 {
+  %lds_ptr = getelementptr [64 x float], [64 x float] addrspace(3)* @lds, i32 0, i32 0
+  %val = load float, float addrspace(3)* %lds_ptr
+  ret float %val
+}
+
+define amdgpu_gfx float @simple_lds_recurse(float %arg0) #0 {
+  %lds_ptr = getelementptr [64 x float], [64 x float] addrspace(3)* @lds, i32 0, i32 0
+  %val = load float, float addrspace(3)* %lds_ptr
+  %res = call amdgpu_gfx float @simple_lds_recurse(float %val)
+  ret float %res
+}
+
 attributes #0 = { nounwind }
 
 ; GCN: amdpal.pipelines:
-; GCN-NEXT:   - .registers:      {}
+; GCN-NEXT:  - .registers:
+; SDAG-NEXT:      0x2e12 (COMPUTE_PGM_RSRC1): 0xaf01ca{{$}}
+; SDAG-NEXT:      0x2e13 (COMPUTE_PGM_RSRC2): 0x8001{{$}}
+; GISEL-NEXT:      0x2e12 (COMPUTE_PGM_RSRC1): 0xaf01cf{{$}}
+; GISEL-NEXT:      0x2e13 (COMPUTE_PGM_RSRC2): 0x8001{{$}}
 ; GCN-NEXT:    .shader_functions:
 ; GCN-NEXT:      dynamic_stack:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x10{{$}}
 ; GCN-NEXT:      dynamic_stack_loop:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x10{{$}}
 ; GCN-NEXT:      multiple_stack:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x24{{$}}
 ; GCN-NEXT:      no_stack:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0{{$}}
 ; GCN-NEXT:      no_stack_call:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0{{$}}
 ; GCN-NEXT:      no_stack_extern_call:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x10{{$}}
 ; GCN-NEXT:      no_stack_extern_call_many_args:
 ; SDAG-NEXT:        .stack_frame_size_in_bytes: 0x90{{$}}
 ; GISEL-NEXT:        .stack_frame_size_in_bytes: 0xd0{{$}}
 ; GCN-NEXT:      no_stack_indirect_call:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x10{{$}}
+; GCN-NEXT:      simple_lds:
+; GCN-NEXT:        .stack_frame_size_in_bytes: 0{{$}}
+; GCN-NEXT:      simple_lds_recurse:
+; GCN-NEXT:        .stack_frame_size_in_bytes: 0x10{{$}}
 ; GCN-NEXT:      simple_stack:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x14{{$}}
 ; GCN-NEXT:      simple_stack_call:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x20{{$}}
 ; GCN-NEXT:      simple_stack_extern_call:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x20{{$}}
 ; GCN-NEXT:      simple_stack_indirect_call:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x20{{$}}
 ; GCN-NEXT:      simple_stack_recurse:
 ; GCN-NEXT:        .stack_frame_size_in_bytes: 0x20{{$}}
 ; GCN-NEXT: ...