[llvm] [AMDGPU] Add object linking support for LDS and named barrier lowering in the middle end (PR #191645)

[Shilei Tian via llvm-commits]

================
@@ -910,7 +911,170 @@ class AMDGPULowerModuleLDS {
     return KernelToCreatedDynamicLDS;
   }
 
+  // Per-TU mode for link-time LDS resolution. Instead of computing a global
+  // layout, create per-function LDS struct declarations so the linker can
+  // assign offsets across TUs.
+  bool runOnModuleLinkTime(Module &M) {
+    bool Changed = superAlignLDSGlobals(M);
+    Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M);
+
+    CallGraph CG(M);
+    FunctionVariableMap KernelLDSUses, FunctionLDSUses;
+    getUsesOfLDSByFunction(CG, M, KernelLDSUses, FunctionLDSUses);
+
+    if (KernelLDSUses.empty() && FunctionLDSUses.empty())
+      return Changed;
+
+    FunctionVariableMap AllLDSUses;
+    for (auto &[F, Vars] : KernelLDSUses)
+      AllLDSUses[F].insert(Vars.begin(), Vars.end());
+    for (auto &[F, Vars] : FunctionLDSUses)
+      AllLDSUses[F].insert(Vars.begin(), Vars.end());
+
+    // Named barriers are handled by AMDGPULowerExecSync; filter them out.
+    for (auto &[F, Vars] : AllLDSUses) {
+      SmallVector<GlobalVariable *> Barriers;
+      for (GlobalVariable *V : Vars)
+        if (AMDGPU::isNamedBarrier(*V))
+          Barriers.push_back(V);
+      for (GlobalVariable *V : Barriers)
+        Vars.erase(V);
+    }
+
+    // Build reverse map: LDS variable -> functions that use it.
+    DenseMap<GlobalVariable *, SmallVector<Function *, 4>> VarToFuncs;
+    for (auto &[F, Vars] : AllLDSUses) {
+      for (GlobalVariable *V : Vars)
+        VarToFuncs[V].push_back(F);
+    }
+
+    // A variable is function-scope iff it has local linkage and exactly one
+    // user function. Everything else is global-scope and must remain as a
+    // standalone external declaration so the linker can assign a single shared
+    // offset.
+    DenseSet<GlobalVariable *> GlobalScopeVars;
+    DenseSet<GlobalVariable *> InternalMultiUserVars;
+    for (auto &[V, Funcs] : VarToFuncs) {
+      if (!V->hasLocalLinkage() || Funcs.size() > 1) {
+        GlobalScopeVars.insert(V);
+        if (V->hasLocalLinkage())
+          InternalMultiUserVars.insert(V);
+      }
+    }
+
+    // Wrap function-scope LDS into per-function structs (unchanged logic,
+    // but global-scope variables are excluded from the set).
+    SmallVector<std::pair<Function *, GlobalVariable *>, 4> FuncToLdsStruct;
+    DenseSet<GlobalVariable *> AllReplacedVars;
+    for (auto &KV : AllLDSUses) {
+      Function *F = KV.first;
+      DenseSet<GlobalVariable *> FuncScopeVars;
+      for (GlobalVariable *V : KV.second)
+        if (!GlobalScopeVars.count(V))
+          FuncScopeVars.insert(V);
+
+      if (FuncScopeVars.empty())
+        continue;
+
+      std::string StructName = ("__amdgpu_lds." + F->getName()).str();
----------------
shiltian wrote:

My bad. My initial thought was, the function would have a unique name, but now I realized that it can also be some internal function, so we do need to have unique ID here. Updated, and also added a test. Thanks for catching it.

https://github.com/llvm/llvm-project/pull/191645