[amdgpu][nfc] Allocate kernel-specific LDS struct deterministically

[JonChesterfield]

Summary

A kernel may have an associated struct for laying out LDS variables.
This patch puts that instance, if present, at a deterministic address by
allocating it at the same time as the module scope instance.

This is relatively likely to be where the instance was allocated anyway (~NFC)
but will allow later patches to calculate where a given field can be found,
which means a function which is only reachable from a single kernel will be
able to access a LDS variable with zero overhead. That will be particularly
helpful for applications that instantiate a function template containing LDS
variables once per kernel.

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

123

this is both a negligible optimisation and decouples the layout from whether dynamic lds alignment is specified, which is useful for calculating where the variable is going to be from a context that doesn't know whether dynamic lds is in play

[JonChesterfield]

Hoping this is uncontroversial - got more patches to follow up with if we can land this,

[rampitec]

Any tests?

[rampitec]

In D127052#3561174, @rampitec wrote:

Any tests?

It is not immediately obvious it is NFC, especially around alignment.

[JonChesterfield]

I think the layout this chooses is always one of the layouts that could have been picked by the iteration order before. Dynamic LDS alignment gets set when encountering the dynamic variable which could be after the other two were allocated.

I'll write some IR to asm cases that would previously fail non-deterministically but will now pass reliably.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

Need to be careful about multiple anonymous functions

[scchan]

In D127052#3560815, @JonChesterfield wrote:

Hoping this is uncontroversial - got more patches to follow up with if we can land this,

This patch essentially turns the module scope LDS layout struct into creating a function scope thing for every kernel using global LDS?

[JonChesterfield]

In D127052#3563622, @scchan wrote:

In D127052#3560815, @JonChesterfield wrote:

Hoping this is uncontroversial - got more patches to follow up with if we can land this,

This patch essentially turns the module scope LDS layout struct into creating a function scope thing for every kernel using global LDS?

No. A patch by Stas some months ago adapted the module scope LDS lowering to stash remaining kernel variables in a per-kernel struct.

All this patch does, if that per-kernel struct is present, is put it at a predictable address instead of leaving it to the whims of block iteration order.

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

How so? By the time they're functions in IR they have unique names as far as I know. If they don't, the current lowering is already broken for them as it assumes a struct can be uniquely associated with a function by deriving the variable name from the function name

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

They don't have to have a name at all. If you have kernels "@0" and "@1" they'll get the same thing here

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

That's exciting. Broken at present, AMDGPULowerModuleLDSPass creates the kernel-specific variable by passing Twine("llvm.amdgcn.kernel.") + F->getName() + ".lds" to StructType::Create() with a suffix ".t".

I'd be inclined to say the right fix for that is to hard error on them on the basis that anonymous kernels are difficult to invoke by name, but failing that we could assign them names in LowerModuleLDS.

Or we could stop using string equivalence to bind a function and a struct together and do something along the lines of metadata instead which should be more robust to things like renaming the kernel.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

You could do what the final emission does, and use the name out of Mangler::getNameWithPrefix

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

Can I postpone fixing that until after this patch, given it's already broken independent of this patch?

[JonChesterfield]

• rebase

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

90

Mangler::getNameWithPrefix is a local state thing - different instances of Mangler will assign different values. We could value::setName the anonymous function but that'll be a change visible outside of LDS and feels inherently wrong.

[JonChesterfield]

Moved the anonymous function problem to D134741, this patch neither creates nor addresses it.

[JonChesterfield]

• Wrote test case, discovered control flow can be simpler, updated implementation

[arsenm]

llvm/test/CodeGen/AMDGPU/lds-frame-extern.ll

2 ↗

(On Diff #463252)

Should use -mtriple=amdgcn-amd-amdhsa, the checks are using the old style mesa relocations for stack

3–5 ↗

(On Diff #463252)

Is there a real reason to test this with all the targets?

[JonChesterfield]

llvm/test/CodeGen/AMDGPU/lds-frame-extern.ll

2 ↗

(On Diff #463252)

Nice, thanks. Dropping to gfx10 only as well, makes the test easier to read.

[JonChesterfield]

• triple to hsa, drop to single target

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

bool AMDGPUCallLowering::lowerFormalArgumentsKernel(
   MachineFunction &MF = B.getMF();
   const GCNSubtarget *Subtarget = &MF.getSubtarget<GCNSubtarget>();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
   const SIRegisterInfo *TRI = Subtarget->getRegisterInfo();
   const SITargetLowering &TLI = *getTLI<SITargetLowering>();
   const DataLayout &DL = F.getParent()->getDataLayout();
 
-  Info->allocateModuleLDSGlobal(F);
+  Info->allocateKnownAddressLDSGlobal(F);
 
   SmallVector<CCValAssign, 16> ArgLocs;
   CCState CCInfo(F.getCallingConv(), F.isVarArg(), MF, ArgLocs, F.getContext());
 
   allocateHSAUserSGPRs(CCInfo, B, MF, *TRI, *Info);
 
   unsigned i = 0;
   const Align KernArgBaseAlign(16);
bool AMDGPUCallLowering::lowerFormalArguments(
   MachineFunction &MF = B.getMF();
   MachineBasicBlock &MBB = B.getMBB();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
   const GCNSubtarget &Subtarget = MF.getSubtarget<GCNSubtarget>();
   const SIRegisterInfo *TRI = Subtarget.getRegisterInfo();
   const DataLayout &DL = F.getParent()->getDataLayout();
 
-  Info->allocateModuleLDSGlobal(F);
+  Info->allocateKnownAddressLDSGlobal(F);
 
   SmallVector<CCValAssign, 16> ArgLocs;
   CCState CCInfo(CC, F.isVarArg(), MF, ArgLocs, F.getContext());
 
   if (Info->hasImplicitBufferPtr()) {
     Register ImplicitBufferPtrReg = Info->addImplicitBufferPtr(*TRI);
     MF.addLiveIn(ImplicitBufferPtrReg, &AMDGPU::SGPR_64RegClass);
     CCInfo.AllocateReg(ImplicitBufferPtrReg);

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

public:
   bool isMemoryBound() const {
     return MemoryBound;
   }
 
   bool needsWaveLimiter() const {
     return WaveLimiter;
   }
 
-  unsigned allocateLDSGlobal(const DataLayout &DL, const GlobalVariable &GV);
-  void allocateModuleLDSGlobal(const Function &F);
+  unsigned allocateLDSGlobal(const DataLayout &DL, const GlobalVariable &GV) {
+    return allocateLDSGlobal(DL, GV, DynLDSAlign);
+  }
+  unsigned allocateLDSGlobal(const DataLayout &DL, const GlobalVariable &GV,
+                             Align Trailing);
+
+  void allocateKnownAddressLDSGlobal(const Function &F);
+
+  // A kernel function may have an associated LDS allocation, and a kernel-scope
+  // LDS allocation must have an associated kernel function
+  static const GlobalVariable *
+  getKernelLDSGlobalFromFunction(const Function &F);
 
   Align getDynLDSAlign() const { return DynLDSAlign; }
 
   void setDynLDSAlign(const DataLayout &DL, const GlobalVariable &GV);
 };
 
 }
 #endif

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

AMDGPUMachineFunction::AMDGPUMachineFunction(const MachineFunction &MF)
   StaticGDSSize = GDSSize;
 
   CallingConv::ID CC = F.getCallingConv();
   if (CC == CallingConv::AMDGPU_KERNEL || CC == CallingConv::SPIR_KERNEL)
     ExplicitKernArgSize = ST.getExplicitKernArgSize(F, MaxKernArgAlign);
 }
 
 unsigned AMDGPUMachineFunction::allocateLDSGlobal(const DataLayout &DL,
-                                                  const GlobalVariable &GV) {
+                                                  const GlobalVariable &GV,
+                                                  Align Trailing) {
   auto Entry = LocalMemoryObjects.insert(std::make_pair(&GV, 0));
   if (!Entry.second)
     return Entry.first->second;
 
   Align Alignment =
       DL.getValueOrABITypeAlignment(GV.getAlign(), GV.getValueType());
 
   unsigned Offset;
   if (GV.getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
     /// TODO: We should sort these to minimize wasted space due to alignment
     /// padding. Currently the padding is decided by the first encountered use
     /// during lowering.
     Offset = StaticLDSSize = alignTo(StaticLDSSize, Alignment);
 
     StaticLDSSize += DL.getTypeAllocSize(GV.getValueType());
 
-    // Update the LDS size considering the padding to align the dynamic shared
-    // memory.
-    LDSSize = alignTo(StaticLDSSize, DynLDSAlign);
+    // Align LDS size to trailing, e.g. for aligning dynamic shared memory
+    LDSSize = alignTo(StaticLDSSize, Trailing);
   } else {
     assert(GV.getAddressSpace() == AMDGPUAS::REGION_ADDRESS &&
            "expected region address space");
 
     Offset = StaticGDSSize = alignTo(StaticGDSSize, Alignment);
     StaticGDSSize += DL.getTypeAllocSize(GV.getValueType());
 
     // FIXME: Apply alignment of dynamic GDS
     GDSSize = StaticGDSSize;
   }
 
   Entry.first->second = Offset;
   return Offset;
 }
 
+const GlobalVariable *
+AMDGPUMachineFunction::getKernelLDSGlobalFromFunction(const Function &F) {
+  const Module *M = F.getParent();
+  std::string KernelLDSName = "llvm.amdgcn.kernel.";
+  KernelLDSName += F.getName();
+  KernelLDSName += ".lds";
+  return M->getNamedGlobal(KernelLDSName);
+}
+
 // This kernel calls no functions that require the module lds struct
 static bool canElideModuleLDS(const Function &F) {
   return F.hasFnAttribute("amdgpu-elide-module-lds");
 }
 
-void AMDGPUMachineFunction::allocateModuleLDSGlobal(const Function &F) {
+void AMDGPUMachineFunction::allocateKnownAddressLDSGlobal(const Function &F) {
   const Module *M = F.getParent();
   if (isModuleEntryFunction()) {
+
+    // Pointer values start from zero, memory allocated per-kernel-launch
+    // Variables can be grouped into a module level struct and a struct per
+    // kernel function by AMDGPULowerModuleLDSPass. If that is done, they
+    // are allocated at statically computable addresses here.
+    //
+    // Address 0
+    // {
+    //   llvm.amdgcn.module.lds
+    // }
+    // alignment padding
+    // {
+    //   llvm.amdgcn.kernel.some-name.lds
+    // }
+    // dynamic lds alignment padding
+
     const GlobalVariable *GV = M->getNamedGlobal("llvm.amdgcn.module.lds");
+    const GlobalVariable *KV = getKernelLDSGlobalFromFunction(F);
+
+    // Avoid allocating dynamic lds padding between these variables
+    Align ModuleLDSTrailing = KV ? Align() : getDynLDSAlign();
+
     if (GV && !canElideModuleLDS(F)) {
-      unsigned Offset = allocateLDSGlobal(M->getDataLayout(), *GV);
+      unsigned Offset =
+          allocateLDSGlobal(M->getDataLayout(), *GV, ModuleLDSTrailing);
       (void)Offset;
       assert(Offset == 0 &&
              "Module LDS expected to be allocated before other LDS");
     }
+
+    if (KV) {
+      unsigned Offset =
+          allocateLDSGlobal(M->getDataLayout(), *KV, getDynLDSAlign());
+      (void)Offset;
+    }
   }
 }
 
 void AMDGPUMachineFunction::setDynLDSAlign(const DataLayout &DL,
                                            const GlobalVariable &GV) {
   assert(DL.getTypeAllocSize(GV.getValueType()).isZero());
 
   Align Alignment =
       DL.getValueOrABITypeAlignment(GV.getAlign(), GV.getValueType());
   if (Alignment <= DynLDSAlign)
     return;
 
   LDSSize = alignTo(StaticLDSSize, Alignment);
   DynLDSAlign = Alignment;
 }

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::LowerFormalArguments(
 
   if (Subtarget->isAmdHsaOS() && AMDGPU::isGraphics(CallConv)) {
     DiagnosticInfoUnsupported NoGraphicsHSA(
         Fn, "unsupported non-compute shaders with HSA", DL.getDebugLoc());
     DAG.getContext()->diagnose(NoGraphicsHSA);
     return DAG.getEntryNode();
   }
 
-  Info->allocateModuleLDSGlobal(Fn);
+  Info->allocateKnownAddressLDSGlobal(Fn);
 
   SmallVector<ISD::InputArg, 16> Splits;
   SmallVector<CCValAssign, 16> ArgLocs;
   BitVector Skipped(Ins.size());
   CCState CCInfo(CallConv, isVarArg, DAG.getMachineFunction(), ArgLocs,
                  *DAG.getContext());
 
   bool IsGraphics = AMDGPU::isGraphics(CallConv);