[amdgpu] Implement dynamic LDS accesses from non-kernel functions

[JonChesterfield]

Summary

The premise here is to allow non-kernel functions to locate external LDS variables without using LDS or extra magic SGPRs to do so.

1/ First it crawls the callgraph to work out which external LDS variables are reachable from a given kernel
2/ Then it creates a new extern char[0] variable for each kernel, which will alias all the other extern LDS variables because that's the documented behaviour of these variables
3/ The address of that variable is written to a lookup table. The global variable is tagged with metadata to track what address it was allocated at by codegen
4/ The assembler builds the lookup table using the metadata
5/ Any non-kernel functions use the same magic intrinsic used by table lookups of non-dynamic LDS variables to find the address to use

Heavy overlap with the code paths taken for other lowering, in particular the same intrinsic is used to pass the dynamic scope information through the same sgpr as for table lookups of static LDS.

[JonChesterfield]

• delete a print statement

[JonChesterfield]

Needs to be rebased on D144221

llvm/lib/Target/AMDGPU/AMDGPUMachineFunction.cpp

214

@arsenm this hooks the same absolute_symbol machinery implemented in D144221 to allow building a table out of dynamic lds addresses in the assembler

[JonChesterfield]

Looking at AMDGPUPromoteAlloca.cpp, the spilling to LDS is disabled for any function which accesses any global that represents dynamic LDS.

This approach will leave a use of dynamic LDS in every kernel that calls any function that accesses dynamic LDS. That will disable spilling for that function. Thus it will be possible to accurately predict where dynamic LDS will start for that kernel.

For non-kernel functions, spilling to LDS is presently disabled as it used to be unreachable from functions.

Thus as currently implemented, dynamic LDS and spilling to LDS are disjoint. We can emit the metadata from IR and check it in codegen, very like in D144221.

We are probably leaving some performance on the table - functions that are only reachable from one kernel could spill to LDS for free - but that seems OK for now.

[JonChesterfield]

• with frame setup, without tests, rebased on a refactored main

[JonChesterfield]

• start testing, fix an inversion bug caught

[JonChesterfield]

• disable the always inlining pass while testing

[JonChesterfield]

• fixes to frame lowering

[JonChesterfield]

• fixes to frame lowering

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPUAlwaysInlinePass.cpp

130 ↗

(On Diff #507935)

This change (and the associated test change) are not intended to land as part of this revision

[JonChesterfield]

Passed internal CI with the always inliner patched (so that tests actually go down this code path) at 829538

[JonChesterfield]

• rebase

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

187

I think this is call is only adding confusion. Just GV->getAddressSpace() == local && !GV->hasInitializer?

I also don't think checking the initializer is quite right. AMDGPUPromoteAlloca checks for GV->hasExternalLinkage() && AllocSize == 0.

[JonChesterfield]

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

187

Thanks! There's some dead code in utils to remove as well, fixing

[JonChesterfield]

• move isDynamicLDS into util, make consistent with promotealloca

[JonChesterfield]

• move isDynamicLDS into util, make consistent with promotealloca

[JonChesterfield]

lds-frame-extern.ll test interacts badly with the always inline pass which doesn't appear to have a hook to disable it. Considering what to do about that, would rather separate the inliner change from the additional lowering, might move that test case change into a separate patch.

It might be better to move the frame lowering code around separate to this patch, will see how that looks

[arsenm]

LGTM with nits

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

869

Typo ‘dyanmically’

935

The GlobalVariable constructor uses Twine, you don’t need to convert it to std::string here

[JonChesterfield]

• move isDynamicLDS into util, make consistent with promotealloca
• review comments

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);
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/AMDGPULegalizerInfo.cpp

if (AS == AMDGPUAS::LOCAL_ADDRESS && GV->hasExternalLinkage()) {
       Type *Ty = GV->getValueType();
       // HIP uses an unsized array `extern __shared__ T s[]` or similar
       // zero-sized type in other languages to declare the dynamic shared
       // memory which size is not known at the compile time. They will be
       // allocated by the runtime and placed directly after the static
       // allocated ones. They all share the same offset.
       if (B.getDataLayout().getTypeAllocSize(Ty).isZero()) {
         // Adjust alignment for that dynamic shared memory array.
-        MFI->setDynLDSAlign(B.getDataLayout(), *cast<GlobalVariable>(GV));
+        Function &F = MF.getFunction();
+        MFI->setDynLDSAlign(F.getParent(), F, *cast<GlobalVariable>(GV));
         LLT S32 = LLT::scalar(32);
         auto Sz =
             B.buildIntrinsic(Intrinsic::amdgcn_groupstaticsize, {S32}, false);
         B.buildIntToPtr(DstReg, Sz);
         MI.eraseFromParent();
         return true;
       }
     }

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

bool isKernelLDS(const Function *F) {
   // isModuleEntryFunction which returns true for more calling conventions
   // than AMDGPU::isKernel does. There's a FIXME on AMDGPU::isKernel.
   // There's also a test that checks that the LDS lowering does not hit on
   // a graphics shader, denoted amdgpu_ps, so stay with the limited case.
   // Putting LDS in the name of the function to draw attention to this.
   return AMDGPU::isKernel(F->getCallingConv());
 }
 
+bool isDynamicLDS(const GlobalVariable *GV) {
+  if (AMDGPU::isLDSVariableToLower(*GV)) {
+    if (!GV->hasInitializer()) {
+      // addrspace(3) without initializer implies cuda/hip extern __shared__
+      // the semantics for such a variable appears to be that all extern
+      // __shared__ variables alias one another
+      return true;
+    }
+  }
+  return false;
+}
+
 class AMDGPULowerModuleLDS : public ModulePass {
 
   static void
   removeLocalVarsFromUsedLists(Module &M,
                                const DenseSet<GlobalVariable *> &LocalVars) {
     // The verifier rejects used lists containing an inttoptr of a constant
     // so remove the variables from these lists before replaceAllUsesWith
     SmallPtrSet<Constant *, 8> LocalVarsSet;
static void markUsedByKernel(IRBuilder<> &Builder, Function *Func,
 
     // An operand bundle on llvm.donothing works because the call instruction
     // survives until after the last pass that needs to account for LDS. It is
     // better than inline asm as the latter survives until the end of codegen. A
     // totally robust solution would be a function with the same semantics as
     // llvm.donothing that takes a pointer to the instance and is lowered to a
     // no-op after LDS is allocated, but that is not presently necessary.
 
+    // The intrinsic is eliminated shortly before instruction selection.
+
     LLVMContext &Ctx = Func->getContext();
 
     Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI());
 
     FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {});
 
     Function *Decl =
         Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {});
public:
   void replaceUsesInInstructionsWithTableLookup(
       Module &M, ArrayRef<GlobalVariable *> ModuleScopeVariables,
       GlobalVariable *LookupTable) {
 
     LLVMContext &Ctx = M.getContext();
     IRBuilder<> Builder(Ctx);
     Type *I32 = Type::getInt32Ty(Ctx);
 
-    // Accesses from a function use the amdgcn_lds_kernel_id intrinsic which
-    // lowers to a read from a live in register. Emit it once in the entry
-    // block to spare deduplicating it later.
-
-    DenseMap<Function *, Value *> tableKernelIndexCache;
-    auto getTableKernelIndex = [&](Function *F) -> Value * {
-      if (tableKernelIndexCache.count(F) == 0) {
-        LLVMContext &Ctx = M.getContext();
-        FunctionType *FTy = FunctionType::get(Type::getInt32Ty(Ctx), {});
-        Function *Decl =
-            Intrinsic::getDeclaration(&M, Intrinsic::amdgcn_lds_kernel_id, {});
-
-        BasicBlock::iterator it =
-            F->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
-        Instruction &i = *it;
-        Builder.SetInsertPoint(&i);
-
-        tableKernelIndexCache[F] = Builder.CreateCall(FTy, Decl, {});
-      }
-
-      return tableKernelIndexCache[F];
-    };
 
     for (size_t Index = 0; Index < ModuleScopeVariables.size(); Index++) {
       auto *GV = ModuleScopeVariables[Index];
 
       for (Use &U : make_early_inc_range(GV->uses())) {
         auto *I = dyn_cast<Instruction>(U.getUser());
         if (!I)
           continue;
 
-        Value *tableKernelIndex = getTableKernelIndex(I->getFunction());
+        Value *tableKernelIndex =
+            getTableLookupKernelIndex(M, I->getFunction());
 
         // So if the phi uses this value multiple times, what does this look
         // like?
         if (auto *Phi = dyn_cast<PHINode>(I)) {
           BasicBlock *BB = Phi->getIncomingBlock(U);
           Builder.SetInsertPoint(&(*(BB->getFirstInsertionPt())));
         } else {
           Builder.SetInsertPoint(I);
for (auto &K : LDSVars) {
                       DL.getTypeAllocSize(GV->getValueType()).getFixedValue());
       if (MostUsed < Candidate)
         MostUsed = Candidate;
     }
 
     return MostUsed.GV;
   }
 
+  DenseMap<Function *, Value *> tableKernelIndexCache;
+  Value *getTableLookupKernelIndex(Module &M, Function *F) {
+    // Accesses from a function use the amdgcn_lds_kernel_id intrinsic which
+    // lowers to a read from a live in register. Emit it once in the entry
+    // block to spare deduplicating it later.
+    if (tableKernelIndexCache.count(F) == 0) {
+      LLVMContext &Ctx = M.getContext();
+      IRBuilder<> Builder(Ctx);
+      FunctionType *FTy = FunctionType::get(Type::getInt32Ty(Ctx), {});
+      Function *Decl =
+          Intrinsic::getDeclaration(&M, Intrinsic::amdgcn_lds_kernel_id, {});
+
+      BasicBlock::iterator it =
+          F->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
+      Instruction &i = *it;
+      Builder.SetInsertPoint(&i);
+
+      tableKernelIndexCache[F] = Builder.CreateCall(FTy, Decl, {});
+    }
+
+    return tableKernelIndexCache[F];
+  }
+
   bool runOnModule(Module &M) override {
     LLVMContext &Ctx = M.getContext();
     CallGraph CG = CallGraph(M);
     bool Changed = superAlignLDSGlobals(M);
 
     Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M);
 
     Changed = true; // todo: narrow this down
for (auto &K : LDSUsesInfo.indirect_access) {
         LDSToKernelsThatNeedToAccessItIndirectly[GV].insert(F);
       }
     }
 
     // Partition variables into the different strategies
     DenseSet<GlobalVariable *> ModuleScopeVariables;
     DenseSet<GlobalVariable *> TableLookupVariables;
     DenseSet<GlobalVariable *> KernelAccessVariables;
+    DenseSet<GlobalVariable *> DynamicVariables;
 
     {
       GlobalVariable *HybridModuleRoot =
           LoweringKindLoc != LoweringKind::hybrid
               ? nullptr
               : chooseBestVariableForModuleStrategy(
                     M.getDataLayout(),
                     LDSToKernelsThatNeedToAccessItIndirectly);
 
       DenseSet<Function *> const EmptySet;
       DenseSet<Function *> const &HybridModuleRootKernels =
           HybridModuleRoot
               ? LDSToKernelsThatNeedToAccessItIndirectly[HybridModuleRoot]
               : EmptySet;
 
       for (auto &K : LDSToKernelsThatNeedToAccessItIndirectly) {
         // Each iteration of this loop assigns exactly one global variable to
         // exactly one of the implementation strategies.
 
         GlobalVariable *GV = K.first;
         assert(AMDGPU::isLDSVariableToLower(*GV));
         assert(K.second.size() != 0);
 
+        if (isDynamicLDS(GV)) {
+          DynamicVariables.insert(GV);
+          continue;
+        }
+
         switch (LoweringKindLoc) {
         case LoweringKind::module:
           ModuleScopeVariables.insert(GV);
           break;
 
         case LoweringKind::table:
           TableLookupVariables.insert(GV);
           break;
DenseSet<GlobalVariable *> DynamicVariables;
             TableLookupVariables.insert(GV);
           }
           break;
         }
         }
       }
 
       assert(ModuleScopeVariables.size() + TableLookupVariables.size() +
-                 KernelAccessVariables.size() ==
+                 KernelAccessVariables.size() + DynamicVariables.size() ==
              LDSToKernelsThatNeedToAccessItIndirectly.size());
-    } // Variables have now been partitioned into the three lowering strategies.
+    } // Variables have now been partitioned into the distinct lowering
+      // strategies.
 
     // If the kernel accesses a variable that is going to be stored in the
     // module instance through a call then that kernel needs to allocate the
     // module instance
     DenseSet<Function *> KernelsThatAllocateModuleLDS =
         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
                                                         ModuleScopeVariables);
     DenseSet<Function *> KernelsThatAllocateTableLDS =
         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
                                                         TableLookupVariables);
+    DenseSet<Function *> KernelsThatAllocateDynamicLDS =
+        kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
+                                                        DynamicVariables);
 
     if (!ModuleScopeVariables.empty()) {
       LDSVariableReplacement ModuleScopeReplacement =
           createLDSVariableReplacement(M, "llvm.amdgcn.module.lds",
                                        ModuleScopeVariables);
 
       appendToCompilerUsed(M,
                            {static_cast<GlobalValue *>(
if (!ModuleScopeVariables.empty()) {
           markUsedByKernel(Builder, &Func, ModuleScopeReplacement.SGV);
 
         } else {
           Func.addFnAttr("amdgpu-elide-module-lds");
         }
       }
     }
 
-    // Create a struct for each kernel for the non-module-scope variables
+    // Create a struct for each kernel for the non-module-scope, non-dynamic
+    // variables
     DenseMap<Function *, LDSVariableReplacement> KernelToReplacement;
     for (Function &Func : M.functions()) {
       if (Func.isDeclaration() || !isKernelLDS(&Func))
         continue;
 
       DenseSet<GlobalVariable *> KernelUsedVariables;
       for (auto &v : LDSUsesInfo.direct_access[&Func]) {
+        if (!isDynamicLDS(v))
           KernelUsedVariables.insert(v);
       }
       for (auto &v : LDSUsesInfo.indirect_access[&Func]) {
+        if (!isDynamicLDS(v))
           KernelUsedVariables.insert(v);
       }
 
       // Variables allocated in module lds must all resolve to that struct,
       // not to the per-kernel instance.
       if (KernelsThatAllocateModuleLDS.contains(&Func)) {
         for (GlobalVariable *v : ModuleScopeVariables) {
           KernelUsedVariables.erase(v);
         }
       }
 
       if (KernelUsedVariables.empty()) {
         // Either used no LDS, or all the LDS it used was also in module
         continue;
       }
 
       // The association between kernel function and LDS struct is done by
       // symbol name, which only works if the function in question has a
       // name This is not expected to be a problem in practice as kernels
       // are called by name making anonymous ones (which are named by the
       // backend) difficult to use. This does mean that llvm test cases need
       // to name the kernels.
for (auto &GV : KernelAccessVariables) {
       DenseSet<GlobalVariable *> Vec;
       Vec.insert(GV);
 
       replaceLDSVariablesWithStruct(M, Vec, Replacement, [](Use &U) {
                                                            return isa<Instruction>(U.getUser());
       });
     }
 
-    if (!KernelsThatAllocateTableLDS.empty()) {
-      // Collect the kernels that allocate table lookup LDS
-      std::vector<Function *> OrderedKernels;
-      {
-        for (Function &Func : M.functions()) {
-          if (Func.isDeclaration())
-            continue;
-          if (!isKernelLDS(&Func))
-            continue;
+    // Collect kernels that allocate table lookup LDS or need to allocate
+    // dynamic LDS
+    if (!KernelsThatAllocateTableLDS.empty() ||
+        !KernelsThatAllocateDynamicLDS.empty()) {
+
 
-          if (KernelsThatAllocateTableLDS.contains(&Func)) {
-            assert(Func.hasName()); // else fatal error earlier
-            OrderedKernels.push_back(&Func);
-          }
-        }
+        if (KernelsThatAllocateTableLDS.contains(&Func) ||
+            KernelsThatAllocateDynamicLDS.contains(&Func)) {
+          assert(Func.hasName()); // else fatal error earlier (todo: make this
+                                  // true)
 
-        // Put them in an arbitrary but reproducible order
-        llvm::sort(OrderedKernels.begin(), OrderedKernels.end(),
-                   [](const Function *lhs, const Function *rhs) -> bool {
-                     return lhs->getName() < rhs->getName();
-                   });
 
-        // Annotate the kernels with their order in this vector
-        LLVMContext &Ctx = M.getContext();
-        IRBuilder<> Builder(Ctx);
 
-        if (OrderedKernels.size() > UINT32_MAX) {
-          // 32 bit keeps it in one SGPR. > 2**32 kernels won't fit on the GPU
-          report_fatal_error("Unimplemented LDS lowering for > 2**32 kernels");
-        }
 
-        for (size_t i = 0; i < OrderedKernels.size(); i++) {
-          Metadata *AttrMDArgs[1] = {
-              ConstantAsMetadata::get(Builder.getInt32(i)),
-          };
-          OrderedKernels[i]->setMetadata("llvm.amdgcn.lds.kernel.id",
-                                         MDNode::get(Ctx, AttrMDArgs));
+      }
+    }
+
+    if (!KernelsThatAllocateTableLDS.empty()) {
+      {
+        LLVMContext &Ctx = M.getContext();
+        IRBuilder<> Builder(Ctx);
 
+        for (size_t i = 0; i < OrderedKernels.size(); i++) {
           markUsedByKernel(Builder, OrderedKernels[i],
                            KernelToReplacement[OrderedKernels[i]].SGV);
         }
       }
 
       // The order must be consistent between lookup table and accesses to
       // lookup table
       std::vector<GlobalVariable *> TableLookupVariablesOrdered(
           TableLookupVariables.begin(), TableLookupVariables.end());
       llvm::sort(TableLookupVariablesOrdered.begin(),
                  TableLookupVariablesOrdered.end(),
                  [](const GlobalVariable *lhs, const GlobalVariable *rhs) {
                    return lhs->getName() < rhs->getName();
                  });
 
+      // element[k] of lookup table is an array of i32 of length
+      // TableLookupVariablesOrdered.size() containing the address of the ith
+      // variable at index i
+
       GlobalVariable *LookupTable = buildLookupTable(
           M, TableLookupVariablesOrdered, OrderedKernels, KernelToReplacement);
       replaceUsesInInstructionsWithTableLookup(M, TableLookupVariablesOrdered,
                                                LookupTable);
     }
 
+    if (!KernelsThatAllocateDynamicLDS.empty()) {
+      LLVMContext &Ctx = M.getContext();
+      IRBuilder<> Builder(Ctx);
+      const DataLayout &DL = M.getDataLayout();
+      Type *I32 = Type::getInt32Ty(Ctx);
+
+      std::vector<Constant *> newDynamicLDS;
+
+      for (auto &func : OrderedKernels) {
+
+        if (KernelsThatAllocateDynamicLDS.contains(func)) {
+
+          Align MaxDynamicAlignment(1);
+          fprintf(stdout, "kernel %s\n", func->getName().str().c_str());
+          for (GlobalVariable *GV : LDSUsesInfo.indirect_access[func]) {
+            if (!DynamicVariables.contains(GV))
+              continue;
+
+            fprintf(stdout, "  uses dynamic variable\n  ");
+            GV->dump();
+
+            MaxDynamicAlignment =
+                std::max(MaxDynamicAlignment, AMDGPU::getAlign(DL, GV));
+          }
+
+          auto emptyCharArray = ArrayType::get(Type::getInt8Ty(Ctx), 0);
+          std::string VarName =
+              Twine("llvm.amdgcn." + func->getName() + ".dynlds").str();
+          GlobalVariable *N = new GlobalVariable(
+              M, emptyCharArray, false, GlobalValue::ExternalLinkage, nullptr,
+              VarName, nullptr, GlobalValue::NotThreadLocal,
+              AMDGPUAS::LOCAL_ADDRESS, false);
+          N->setAlignment(MaxDynamicAlignment);
+
+          markUsedByKernel(Builder, func, N);
+
+          assert(isDynamicLDS(N));
+
+          auto GEP = ConstantExpr::getGetElementPtr(
+              emptyCharArray, N, ConstantInt::get(I32, 0), true);
+          newDynamicLDS.push_back(ConstantExpr::getPtrToInt(GEP, I32));
+        } else {
+          newDynamicLDS.push_back(PoisonValue::get(I32));
+        }
+      }
+
+      ArrayType *t = ArrayType::get(I32, OrderedKernels.size());
+      Constant *init = ConstantArray::get(t, newDynamicLDS);
+      GlobalVariable *table = new GlobalVariable(
+          M, t, true, GlobalValue::InternalLinkage, init,
+          "llvm.amdgcn.dynlds.offset.table", nullptr,
+          GlobalValue::NotThreadLocal, AMDGPUAS::CONSTANT_ADDRESS);
+
+      for (GlobalVariable *GV : DynamicVariables) {
+        for (Use &U : make_early_inc_range(GV->uses())) {
+          auto *I = dyn_cast<Instruction>(U.getUser());
+          if (!I)
+            continue;
+          if (isKernelLDS(I->getFunction()))
+            continue;
+
+          Value *tableKernelIndex =
+              getTableLookupKernelIndex(M, I->getFunction());
+          if (auto *Phi = dyn_cast<PHINode>(I)) {
+            BasicBlock *BB = Phi->getIncomingBlock(U);
+            Builder.SetInsertPoint(&(*(BB->getFirstInsertionPt())));
+          } else {
+            Builder.SetInsertPoint(I);
+          }
+
+          Value *GEPIdx[2] = {
+              ConstantInt::get(I32, 0),
+              tableKernelIndex,
+          };
+
+          Value *Address = Builder.CreateInBoundsGEP(
+              table->getValueType(), table, GEPIdx, GV->getName());
+
+          Value *loaded = Builder.CreateLoad(I32, Address);
+
+          Value *replacement =
+              Builder.CreateIntToPtr(loaded, GV->getType(), GV->getName());
+
+          U.set(replacement);
+        }
+      }
+    }
+
     for (auto &GV : make_early_inc_range(M.globals()))
       if (AMDGPU::isLDSVariableToLower(GV)) {
-
         // probably want to remove from used lists
         GV.removeDeadConstantUsers();
         if (GV.use_empty())
           GV.eraseFromParent();
       }
 
     return Changed;
   }

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 GlobalVariable *getKernelDynLDSGlobalFromFunction(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);
+  void setDynLDSAlign(Module *M, const Function &F, 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);
 }
 
+GlobalVariable *
+AMDGPUMachineFunction::getKernelDynLDSGlobalFromFunction(Module *M,
+                                                         const Function &F) {
+  std::string KernelDynLDSName = "llvm.amdgcn.";
+  KernelDynLDSName += F.getName();
+  KernelDynLDSName += ".dynlds";
+  return M->getNamedGlobal(KernelDynLDSName);
+}
+
 // 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;
-  }
+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))));
 
-  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;
-}
-
-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);
+    GlobalVariable *Dyn = getKernelDynLDSGlobalFromFunction(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);
+    }
+
+    if (Dyn) {
+      setDynLDSAlign(M, F, *Dyn);
     }
   }
 }
 
 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;
       }
     }
   }
   return {};
 }
 
-void AMDGPUMachineFunction::setDynLDSAlign(const DataLayout &DL,
+void AMDGPUMachineFunction::setDynLDSAlign(Module *M, const Function &F,
                                            const GlobalVariable &GV) {
+  const DataLayout &DL = M->getDataLayout();
   assert(DL.getTypeAllocSize(GV.getValueType()).isZero());
 
   Align Alignment =
       DL.getValueOrABITypeAlignment(GV.getAlign(), GV.getValueType());
   if (Alignment <= DynLDSAlign)
     return;
 
   LDSSize = alignTo(StaticLDSSize, Alignment);
   DynLDSAlign = Alignment;
+
+  // If there is a dynamic LDS variable associated with this function, update it
+  // whenever changing the address of dynamic LDS, aka when calling this
+  // function
+  GlobalVariable *Dyn =
+      AMDGPUMachineFunction::getKernelDynLDSGlobalFromFunction(M, F);
+  if (Dyn) {
+    recordLDSAbsoluteAddress(M, Dyn, LDSSize);
+  }
 }

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);
if (GSD->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS &&
       // HIP uses an unsized array `extern __shared__ T s[]` or similar
       // zero-sized type in other languages to declare the dynamic shared
       // memory which size is not known at the compile time. They will be
       // allocated by the runtime and placed directly after the static
       // allocated ones. They all share the same offset.
       if (DAG.getDataLayout().getTypeAllocSize(Ty).isZero()) {
         assert(PtrVT == MVT::i32 && "32-bit pointer is expected.");
         // Adjust alignment for that dynamic shared memory array.
-        MFI->setDynLDSAlign(DAG.getDataLayout(), *cast<GlobalVariable>(GV));
+        Function &F = DAG.getMachineFunction().getFunction();
+        MFI->setDynLDSAlign(F.getParent(), F, *cast<GlobalVariable>(GV));
         return SDValue(
             DAG.getMachineNode(AMDGPU::GET_GROUPSTATICSIZE, DL, PtrVT), 0);
       }
     }
     return AMDGPUTargetLowering::LowerGlobalAddress(MFI, Op, DAG);
   }
 
   if (GSD->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {

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

static bool shouldLowerLDSToStruct(const GlobalVariable &GV,
 
   return Ret;
 }
 
 bool isLDSVariableToLower(const GlobalVariable &GV) {
   if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) {
     return false;
   }
+  if (GV.isConstant()) {
+    // A constant undef variable can't be written to, and any load is
+    // undef, so it should be eliminated by the optimizer. It could be
+    // dropped by the back end if not. This pass skips over it.
+    return false;
+  }
   if (!GV.hasInitializer()) {
     // addrspace(3) without initializer implies cuda/hip extern __shared__
     // the semantics for such a variable appears to be that all extern
-    // __shared__ variables alias one another, in which case this transform
-    // is not required
-    return false;
+    // __shared__ variables alias one another. This hits different handling.
+    return true;
   }
   if (!isa<UndefValue>(GV.getInitializer())) {
     // Initializers are unimplemented for LDS address space.
     // Leave such variables in place for consistent error reporting.
     return false;
   }
-  if (GV.isConstant()) {
-    // A constant undef variable can't be written to, and any load is
-    // undef, so it should be eliminated by the optimizer. It could be
-    // dropped by the back end if not. This pass skips over it.
-    return false;
-  }
   return true;
 }
 
 std::vector<GlobalVariable *> findLDSVariablesToLower(Module &M,
                                                       const Function *F) {
   std::vector<llvm::GlobalVariable *> LocalVars;
   for (auto &GV : M.globals()) {
     if (!isLDSVariableToLower(GV)) {

llvm/test/CodeGen/AMDGPU/lower-module-lds-inactive.ll

 @var2 = addrspace(3) global float undef
 
 ; constant variables are left to the optimizer / error diagnostics
 ; CHECK: @const_undef = addrspace(3) constant i32 undef
 ; CHECK: @const_with_init = addrspace(3) constant i64 8
 @const_undef = addrspace(3) constant i32 undef
 @const_with_init = addrspace(3) constant i64 8
 
-; External and constant are both left to the optimizer / error diagnostics
-; CHECK: @extern = external addrspace(3) global i32
-@extern = external addrspace(3) global i32
-
 ; Use of an addrspace(3) variable with an initializer is skipped,
 ; so as to preserve the unimplemented error from llc
 ; CHECK: @with_init = addrspace(3) global i64 0
 @with_init = addrspace(3) global i64 0
 
 ; Only local addrspace variables are transformed
 ; CHECK: @addr4 = addrspace(4) global i64 undef
 @addr4 = addrspace(4) global i64 undef

 ; CHECK: @llvm.compiler.used = appending global [1 x ptr] [ptr addrspacecast (ptr addrspace(3) @var2 to ptr)], section "llvm.metadata"
 @llvm.compiler.used = appending global [1 x ptr] [ptr addrspacecast (ptr addrspace(3) @var2 to ptr)], section "llvm.metadata"
 
 ; Access from a function would cause lowering for non-excluded cases
 ; CHECK-LABEL: @use_variables()
 ; CHECK: %c0 = load i32, ptr addrspace(3) @const_undef, align 4
 ; CHECK: %c1 = load i64, ptr addrspace(3) @const_with_init, align 4
 ; CHECK: %v0 = atomicrmw add ptr addrspace(3) @with_init, i64 1 seq_cst
-; CHECK: %v1 = cmpxchg ptr addrspace(3) @extern, i32 4, i32 %c0 acq_rel monotonic
-; CHECK: %v2 = atomicrmw add ptr addrspace(4) @addr4, i64 %c1 monotonic
+; CHECK: %v1 = atomicrmw add ptr addrspace(4) @addr4, i64 %c1 monotonic
 define void @use_variables() {
   %c0 = load i32, ptr addrspace(3) @const_undef, align 4
   %c1 = load i64, ptr addrspace(3) @const_with_init, align 4
   %v0 = atomicrmw add ptr addrspace(3) @with_init, i64 1 seq_cst
-  %v1 = cmpxchg ptr addrspace(3) @extern, i32 4, i32 %c0 acq_rel monotonic
-  %v2 = atomicrmw add ptr addrspace(4) @addr4, i64 %c1 monotonic
+  %v1 = atomicrmw add ptr addrspace(4) @addr4, i64 %c1 monotonic
   ret void
 }
 
 ; CHECK-LABEL: @kern_use()
 ; CHECK: %inc = atomicrmw add ptr addrspace(3) @llvm.amdgcn.kernel.kern_use.lds, i32 1 monotonic, align 4
 define amdgpu_kernel void @kern_use() {
   %inc = atomicrmw add ptr addrspace(3) @var1, i32 1 monotonic
   call void @use_variables()
   ret void
 }