[amdgpu][nfc] Replace ad hoc LDS frame recalculation with absolute_symbol MD

[JonChesterfield]

Summary

Post ISel, LDS variables are absolute values. Representing them as
such is simpler than the frame recalculation currently used to build assembler
tables from their addresses.

This is a precursor to lowering dynamic/external LDS accesses from non-kernel
functions.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

515

F.getParent

599

Ditto

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

124

Changing the IR in a codegen pass is bad. Can you set this earlier?

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

515

Nope, F.getParent will yield a const Module, need a non-const one in order to get a non-const GlobalVariable in order to call setMetadata on it. Otherwise would just call F.getParent within allocateKnownAddressLDSGlobal.

516

Making the function argument non-const would be simpler for the allocateKnownAddress call but makes a mess of some of the call sites

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

124

Setting metadata on the IR during codegen is not pretty but it's also inconsequential. The frame calculation this deletes is also quite ugly though.

For the existing cases we can set the metadata in the IR lowering (carefully, but accurately enough). It'll make changes to promote alloca fragile. That won't work for external/dynamic LDS as we don't know that address until the end of ISel.

So yes... if we want a hard no on setting metadata to record information learned during codegen, we can avoid it. Basically by lifting the remaining LDS lowering currently done by codegen wholly out of codegen. That'll change how dynamic LDS is handled from kernels which I was hoping to avoid.

If we're up for tagging variables with metadata at this point we can write the address chosen for dynamic LDS variables the same way and the lowerConstant path added here will do the right thing for them.

[JonChesterfield]

Created D144233 (which was supposed to be a diff relative to this patch but actually has all the line noise in it anyway).
That is not ready to go but shows how function scope external LDS can be layered on the same metadata scheme proposed here. Main difference is that we really don't know that address until codegen time, at least as presently implemented.

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

120

Note to self - this shouldn't be dead after this patch - find out what happened to the lowering path for the currently unused "direct" scheme

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

120

I think a block was lost in a git merge - AMDGPUIselLowering should have a block in it like:

if (G->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
  if (!MFI->isModuleEntryFunction()) {
    if (const GlobalVariable *GVar = dyn_cast<GlobalVariable>(GV)) {
      if (AMDGPUMachineFunction::isKnownAddressLDSGlobal(*GVar)) {
        unsigned Offset =
            AMDGPUMachineFunction::calculateKnownAddressOfLDSGlobal(*GVar);
        return DAG.getConstant(Offset, SDLoc(Op), Op.getValueType());
      }
    }
  }
}

that allows direct access to kernel allocated variables from functions where the access is unambiguous. This highlights missing test coverage for the "kernel" lowering strategy (from IR to ISA). Bug is latent but will complicate getting rid of the calculateKnownAddressOfLDSGlobal function.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

124

So you're implying the promote alloca introduced LDS still has possibly unpredictable offsets assigned for it?

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

120

So if you set the absolute_address in the IR pass, and LowerGlobalAddress respected it, how would that be less reliable?

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

124

Yep. It's a phase ordering problem.

promoteAlloca needs to know how much LDS is already allocated to estimate how much more it can use. Some of that might be allocated for non-kernel functions, so LowerModuleLDS runs before promoteAlloca to make those allocations visible to it.

LowerModuleLDS thus doesn't know how much (if any) extra will be introduced by promoteAlloca.At least, it doesn't at present.

Then the lowered frame looks like:
{

managed by lowermodulelds
introduced by promotealloca (and maybe other things? none I know of)
maybe alignment
dynamic lds goes here

}

Also if promotealloca introduces more than one lds variable (haven't checked, seems reasonable that it might) and introduces them as independent values with different alignment, then we're back to allocate-in-dag-traversal order, in which case the address of dynamic lds is unstable.

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

120

Would work identically for the module/kernel specific structures as they are deterministic.

Actually somewhat prettier to set those via metadata as it avoids the problem of codegen a function before the corresponding kernel. The above ^ would change to look at metadata on the variables.

That will unwind a little if something after lowermodulelds decides to change those structs, but that's manageable.

Even if we can't mutate the IR during codegen for dynamic lds, I think that's an objective improvement to the current lowering scheme involving repeated frame calculations.

[JonChesterfield]

• rebase

[JonChesterfield]

• Set the metadata in lowermodulelds - test noise update still todo

[JonChesterfield]

• factor out the md parse

[JonChesterfield]

• drop one function, move a second out of the header

[JonChesterfield]

This is still a NFC in terms of machine code but introducing the metadata during LDS lowering introduces a lot of churn into the tests, leaving updating those for a later patch.

Strategy of emitting metadata in the IR pass and then checking allocation matches it removes the frame calculation functions from AMDGPUMachineFunction and reduces (but doesn't quite eliminate) the function/struct symbol name string effort. It's probably possible to tag functions with the struct directly (using more metadata) to remove that entirely.

The primary drawback is that this uses metadata to convey information where we will miscompile if that metadata is removed, directly at odds with the design intent of metadata. However I claim that ship has sailed, and we already assume that some metadata makes it to the backend safely. That the lowering pass runs as part of codegen, not as part of per-TU opt, reduces the risk of another pass invalidating it.

Aside from the metadata-might-be-invalidated hazard, I consider this approach better than the frame calculation in every respect, and would have done it this way around originally if I'd noticed the absolute symbol construct before writing the frame calculation approach.

This is ready for code review - hopefully nothing too surprising in the implementation - but needs a large mechanical update to eight tests to introduce the metadata and change some integers.

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

112

Function has four call sites across three files, think it's probably worth factoring out.

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

108–116

Might be better to specialise this to take a global, check for absolute_symbol and rename the function to match. optional<uint32> getAbsoluteSymbolMetadata(const GlobalVariable *) perhaps.

[arsenm]

In D144221#4152246, @JonChesterfield wrote:

The primary drawback is that this uses metadata to convey information where we will miscompile if that metadata is removed, directly at odds with the design intent of metadata.

absolute_symbol is one of the few cases for global metadata which you are specifically not allowed to drop as per the LangRef.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

666 ↗

(On Diff #500400)

Don't like variables shadowing type names (also hate the LLVM capitalized naming convention). auto *IntTy?

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

115

Early return (or just move the null check to the caller).

The verifier should really be enforcing the operand count so every user doesn't need to check that it's 1. Should add the verifier check and drop the check

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

668 ↗

(On Diff #500400)

Going based on the langref, this is incorrect usage of absolute_symbol. It seems to expect 2 operands, indicating an absolute range. I assume you can specify specific address, specific address + 1 for a fixed address. Again demonstrates that the verifier really should check this

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

115

Also there is already GlobalValue::getAbsoluteSymbolRange()

[JonChesterfield]

• review comments, update tests

[JonChesterfield]

I think that's all comments fixed except for teaching the verifier about range metadata which should be done separately to this.

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

826 ↗

(On Diff #504003)

This is right but moderately annoying to extend for D144233 - would rather land this as is and then amend

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

175

This should probably be an attribute since it doesn't have the special case must-not-drop property but I guess it already was metadata

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->allocateKnownAddressLDSGlobal(F);
+  Info->allocateKnownAddressLDSGlobal(MF.getFunction().getParent(), 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);
   const unsigned BaseOffset = Subtarget->getExplicitKernelArgOffset(F);
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->allocateKnownAddressLDSGlobal(F);
+  Info->allocateKnownAddressLDSGlobal(MF.getFunction().getParent(), 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/AMDGPUMCInstLower.cpp

 //===- AMDGPUMCInstLower.cpp - Lower AMDGPU MachineInstr to an MCInst -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 /// Code to lower AMDGPU MachineInstrs to their corresponding MCInst.
 //
 //===----------------------------------------------------------------------===//
 //
 
 #include "AMDGPUMCInstLower.h"
+#include "AMDGPU.h"
 #include "AMDGPUAsmPrinter.h"
 #include "AMDGPUMachineFunction.h"
 #include "AMDGPUTargetMachine.h"
 #include "MCTargetDesc/AMDGPUInstPrinter.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/IR/Constants.h"
bool AMDGPUAsmPrinter::lowerOperand(const MachineOperand &MO,
   const GCNSubtarget &STI = MF->getSubtarget<GCNSubtarget>();
   AMDGPUMCInstLower MCInstLowering(OutContext, STI, *this);
   return MCInstLowering.lowerOperand(MO, MCOp);
 }
 
 const MCExpr *AMDGPUAsmPrinter::lowerConstant(const Constant *CV) {
 
   // Intercept LDS variables with known addresses
-  if (const GlobalVariable *GV = dyn_cast<GlobalVariable>(CV)) {
-    if (AMDGPUMachineFunction::isKnownAddressLDSGlobal(*GV)) {
-      unsigned offset =
-          AMDGPUMachineFunction::calculateKnownAddressOfLDSGlobal(*GV);
-      Constant *C = ConstantInt::get(CV->getContext(), APInt(32, offset));
-      return AsmPrinter::lowerConstant(C);
+  if (const GlobalVariable *GV = dyn_cast<const GlobalVariable>(CV)) {
+    if (GV->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
+      auto MD = GV->getMetadata(LLVMContext::MD_absolute_symbol);
+      if (MD && MD->getNumOperands() == 1) {
+        if (ConstantInt *KnownSize =
+                mdconst::extract<ConstantInt>(MD->getOperand(0))) {
+          if (KnownSize->getZExtValue() <= UINT32_MAX) {
+            return AsmPrinter::lowerConstant(KnownSize);
+          }
+        }
+      }
     }
   }
 
   if (const MCExpr *E = lowerAddrSpaceCast(TM, CV, OutContext))
     return E;
   return AsmPrinter::lowerConstant(CV);
 }
 

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.h

public:
 
   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);
+  void allocateKnownAddressLDSGlobal(Module *M, const Function &F);
 
   // A kernel function may have an associated LDS allocation, and a kernel-scope
   // LDS allocation must have an associated kernel function
 
   // LDS allocation should have an associated kernel function
   static const Function *
   getKernelLDSFunctionFromGlobal(const GlobalVariable &GV);
-  static const GlobalVariable *
-  getKernelLDSGlobalFromFunction(const Function &F);
-
-  // Module or kernel scope LDS variable
-  static bool isKnownAddressLDSGlobal(const GlobalVariable &GV);
-  static unsigned calculateKnownAddressOfLDSGlobal(const GlobalVariable &GV);
+  static GlobalVariable *getKernelLDSGlobalFromFunction(Module *M,
+                                                        const Function &F);
 
   static std::optional<uint32_t> getLDSKernelIdMetadata(const Function &F);
 
   Align getDynLDSAlign() const { return DynLDSAlign; }
 
   void setDynLDSAlign(const DataLayout &DL, const GlobalVariable &GV);
 };
 
 }
 #endif

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 "AMDGPU.h"
 #include "AMDGPUPerfHintAnalysis.h"
 #include "AMDGPUSubtarget.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/Target/TargetMachine.h"
 
 using namespace llvm;
 
 AMDGPUMachineFunction::AMDGPUMachineFunction(const Function &F,
                                              const AMDGPUSubtarget &ST)
     : IsEntryFunction(AMDGPU::isEntryFunctionCC(F.getCallingConv())),
       IsModuleEntryFunction(
unsigned AMDGPUMachineFunction::allocateLDSGlobal(const DataLayout &DL,
   }
 
   Entry.first->second = Offset;
   return Offset;
 }
 
 static constexpr StringLiteral ModuleLDSName = "llvm.amdgcn.module.lds";
 
-bool AMDGPUMachineFunction::isKnownAddressLDSGlobal(const GlobalVariable &GV) {
-  auto name = GV.getName();
-  return (name == ModuleLDSName) ||
-         (name.startswith("llvm.amdgcn.kernel.") && name.endswith(".lds"));
-}
-
 const Function *AMDGPUMachineFunction::getKernelLDSFunctionFromGlobal(
     const GlobalVariable &GV) {
   const Module &M = *GV.getParent();
   StringRef N(GV.getName());
   if (N.consume_front("llvm.amdgcn.kernel.") && N.consume_back(".lds")) {
     return M.getFunction(N);
   }
   return nullptr;
 }
 
-const GlobalVariable *
-AMDGPUMachineFunction::getKernelLDSGlobalFromFunction(const Function &F) {
-  const Module *M = F.getParent();
+GlobalVariable *
+AMDGPUMachineFunction::getKernelLDSGlobalFromFunction(Module *M,
+                                                      const Function &F) {
   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");
 }
 
-unsigned AMDGPUMachineFunction::calculateKnownAddressOfLDSGlobal(
-    const GlobalVariable &GV) {
-  // module.lds, then alignment padding, then kernel.lds, then other variables
-  // if any
-
-  assert(isKnownAddressLDSGlobal(GV));
-  unsigned Offset = 0;
-
-  if (GV.getName() == ModuleLDSName) {
-    return 0;
-  }
-
-  const Module *M = GV.getParent();
-  const DataLayout &DL = M->getDataLayout();
-
-  const GlobalVariable *GVM = M->getNamedGlobal(ModuleLDSName);
-  const Function *f = getKernelLDSFunctionFromGlobal(GV);
-
-  // Account for module.lds if allocated for this function
-  if (GVM && f && !canElideModuleLDS(*f)) {
-    // allocator aligns this to var align, but it's zero to begin with
-    Offset += DL.getTypeAllocSize(GVM->getValueType());
-  }
-
-  // No dynamic LDS alignment done by allocateModuleLDSGlobal
-  Offset = alignTo(
-      Offset, DL.getValueOrABITypeAlignment(GV.getAlign(), GV.getValueType()));
-
-  return Offset;
+static void recordLDSAbsoluteAddress(Module *M, GlobalVariable *GV,
+                                     uint32_t Address) {
+  // Write the specified address into metadata where it can be retrieved by the
+  // assembler
+  LLVMContext &Ctx = M->getContext();
+  auto Type = M->getDataLayout().getIntPtrType(Ctx, AMDGPUAS::LOCAL_ADDRESS);
+  GV->setMetadata(LLVMContext::MD_absolute_symbol,
+                  MDNode::get(Ctx, ConstantAsMetadata::get(
+                                       ConstantInt::get(Type, Address))));
 }
 
-void AMDGPUMachineFunction::allocateKnownAddressLDSGlobal(const Function &F) {
-  const Module *M = F.getParent();
+void AMDGPUMachineFunction::allocateKnownAddressLDSGlobal(Module *M,
+                                                          const Function &F) {
 
   // This function is called before allocating any other LDS so that it can
   // reliably put values at known addresses. Consequently, dynamic LDS, if
   // present, will not yet have been allocated
 
   assert(getDynLDSAlign() == Align() && "dynamic LDS not yet allocated");
 
   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
     // }
     // other variables, e.g. dynamic lds, allocated after this call
 
-    const GlobalVariable *GV = M->getNamedGlobal(ModuleLDSName);
-    const GlobalVariable *KV = getKernelLDSGlobalFromFunction(F);
+    GlobalVariable *GV = M->getNamedGlobal(ModuleLDSName);
+    GlobalVariable *KV = getKernelLDSGlobalFromFunction(M, F);
 
     if (GV && !canElideModuleLDS(F)) {
-      assert(isKnownAddressLDSGlobal(*GV));
       unsigned Offset = allocateLDSGlobal(M->getDataLayout(), *GV, Align());
-      (void)Offset;
-      assert(Offset == calculateKnownAddressOfLDSGlobal(*GV) &&
-             "Module LDS expected to be allocated before other LDS");
+      recordLDSAbsoluteAddress(M, GV, Offset);
     }
 
     if (KV) {
       // The per-kernel offset is deterministic because it is allocated
       // before any other non-module LDS variables.
-      assert(isKnownAddressLDSGlobal(*KV));
       unsigned Offset = allocateLDSGlobal(M->getDataLayout(), *KV, Align());
-      (void)Offset;
-      assert(Offset == calculateKnownAddressOfLDSGlobal(*KV) &&
-             "Kernel LDS expected to be immediately after module LDS");
+      recordLDSAbsoluteAddress(M, KV, Offset);
     }
   }
 }
 
 std::optional<uint32_t>
 AMDGPUMachineFunction::getLDSKernelIdMetadata(const Function &F) {
   auto MD = F.getMetadata("llvm.amdgcn.lds.kernel.id");
   if (MD && MD->getNumOperands() == 1) {
     ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(0));
     if (KnownSize) {
       uint64_t V = KnownSize->getZExtValue();
       if (V <= UINT32_MAX) {
         return V;
       }
     }
   }

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->allocateKnownAddressLDSGlobal(Fn);
+  Info->allocateKnownAddressLDSGlobal(MF.getFunction().getParent(), 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);