[AMDGPU] Add dynamic LDS size implicit kernel argument to CO-v5 (#65273)

"hidden_dynamic_lds_size" argument will be added in the reserved section
at offset 120 of the implicit argument layout.
Add "isDynamicLDSUsed" flag to AMDGPUMachineFunction to identify if a
function uses dynamic LDS.

hidden argument will be added in below cases:

- LDS global is used in the kernel.
- Kernel calls a function which uses LDS global.
- LDS pointer is passed as argument to kernel itself.

1 files changed, 30 insertions, 9 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp b/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp
index 323462e..3177729 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp
@@ -19,6 +19,26 @@
 
 using namespace llvm;
 
+static const GlobalVariable *
+getKernelDynLDSGlobalFromFunction(const Function &F) {
+  const Module *M = F.getParent();
+  SmallString<64> KernelDynLDSName("llvm.amdgcn.");
+  KernelDynLDSName += F.getName();
+  KernelDynLDSName += ".dynlds";
+  return M->getNamedGlobal(KernelDynLDSName);
+}
+
+static bool hasLDSKernelArgument(const Function &F) {
+  for (const Argument &Arg : F.args()) {
+    Type *ArgTy = Arg.getType();
+    if (auto PtrTy = dyn_cast<PointerType>(ArgTy)) {
+      if (PtrTy->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
+        return true;
+    }
+  }
+  return false;
+}
+
 AMDGPUMachineFunction::AMDGPUMachineFunction(const Function &F,
                                              const AMDGPUSubtarget &ST)
     : IsEntryFunction(AMDGPU::isEntryFunctionCC(F.getCallingConv())),
@@ -65,6 +85,10 @@ AMDGPUMachineFunction::AMDGPUMachineFunction(const Function &F,
   Attribute NSZAttr = F.getFnAttribute("no-signed-zeros-fp-math");
   NoSignedZerosFPMath =
       NSZAttr.isStringAttribute() && NSZAttr.getValueAsString() == "true";
+
+  const GlobalVariable *DynLdsGlobal = getKernelDynLDSGlobalFromFunction(F);
+  if (DynLdsGlobal || hasLDSKernelArgument(F))
+    UsesDynamicLDS = true;
 }
 
 unsigned AMDGPUMachineFunction::allocateLDSGlobal(const DataLayout &DL,
@@ -139,15 +163,6 @@ unsigned AMDGPUMachineFunction::allocateLDSGlobal(const DataLayout &DL,
   return Offset;
 }
 
-static const GlobalVariable *
-getKernelDynLDSGlobalFromFunction(const Function &F) {
-  const Module *M = F.getParent();
-  std::string KernelDynLDSName = "llvm.amdgcn.";
-  KernelDynLDSName += F.getName();
-  KernelDynLDSName += ".dynlds";
-  return M->getNamedGlobal(KernelDynLDSName);
-}
-
 std::optional<uint32_t>
 AMDGPUMachineFunction::getLDSKernelIdMetadata(const Function &F) {
   // TODO: Would be more consistent with the abs symbols to use a range
@@ -210,3 +225,9 @@ void AMDGPUMachineFunction::setDynLDSAlign(const Function &F,
     }
   }
 }
+
+void AMDGPUMachineFunction::setUsesDynamicLDS(bool DynLDS) {
+  UsesDynamicLDS = DynLDS;
+}
+
+bool AMDGPUMachineFunction::isDynamicLDSUsed() const { return UsesDynamicLDS; }