[AMDGPU] Let LowerModuleLDS run twice on the same module

See llvm#81491
Split from llvm#75333

Author: Pierre-vh

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

@@ -327,16 +327,18 @@ class AMDGPULowerModuleLDS {
     return convertUsersOfConstantsToInstructions(LDSGlobals);
   }
 
+  std::optional<bool> HasAbsoluteGVs;
+
 public:
   AMDGPULowerModuleLDS(const AMDGPUTargetMachine &TM_) : TM(TM_) {}
 
   using FunctionVariableMap = DenseMap<Function *, DenseSet<GlobalVariable *>>;
 
   using VariableFunctionMap = DenseMap<GlobalVariable *, DenseSet<Function *>>;
 
-  static void getUsesOfLDSByFunction(CallGraph const &CG, Module &M,
-                                     FunctionVariableMap &kernels,
-                                     FunctionVariableMap &functions) {
+  void getUsesOfLDSByFunction(CallGraph const &CG, Module &M,
+                              FunctionVariableMap &kernels,
+                              FunctionVariableMap &functions) {
 
     // Get uses from the current function, excluding uses by called functions
     // Two output variables to avoid walking the globals list twice
@@ -345,10 +347,18 @@ class AMDGPULowerModuleLDS {
         continue;
       }
 
-      if (GV.isAbsoluteSymbolRef()) {
-        report_fatal_error(
-            "LDS variables with absolute addresses are unimplemented.");
-      }
+      // Check if the module is consistent: either all GVs are absolute (happens
+      // when we run the pass more than once), or none are.
+      if (HasAbsoluteGVs.has_value()) {
+        if (*HasAbsoluteGVs != GV.isAbsoluteSymbolRef()) {
+          report_fatal_error(
+              "Module cannot mix absolute and non-absolute LDS GVs");
+        }
+      } else
+        HasAbsoluteGVs = GV.isAbsoluteSymbolRef();
+
+      if (GV.isAbsoluteSymbolRef())
+        continue;
 
       for (User *V : GV.users()) {
         if (auto *I = dyn_cast<Instruction>(V)) {
@@ -368,7 +378,7 @@ class AMDGPULowerModuleLDS {
     FunctionVariableMap indirect_access;
   };
 
-  static LDSUsesInfoTy getTransitiveUsesOfLDS(CallGraph const &CG, Module &M) {
+  LDSUsesInfoTy getTransitiveUsesOfLDS(CallGraph const &CG, Module &M) {
 
     FunctionVariableMap direct_map_kernel;
     FunctionVariableMap direct_map_function;

llvm/test/CodeGen/AMDGPU/lds-reject-absolute-addresses.ll renamed to llvm/test/CodeGen/AMDGPU/lds-reject-mixed-absolute-addresses.ll

@@ -2,8 +2,9 @@
 ; RUN: not --crash opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds < %s 2>&1 | FileCheck %s
 
 @var1 = addrspace(3) global i32 undef, !absolute_symbol !0
+@var2 = addrspace(3) global i32 undef
 
-; CHECK: LLVM ERROR: LDS variables with absolute addresses are unimplemented.
+; CHECK: Module cannot mix absolute and non-absolute LDS GVs
 define amdgpu_kernel void @kern() {
   %val0 = load i32, ptr addrspace(3) @var1
   %val1 = add i32 %val0, 4
@@ -12,4 +13,3 @@ define amdgpu_kernel void @kern() {
 }
 
 !0 = !{i32 0, i32 1}
-

llvm/test/CodeGen/AMDGPU/lds-run-twice.ll

@@ -0,0 +1,14 @@
+; RUN: opt -S -mtriple=amdgcn-- -amdgpu-lower-module-lds %s -o %t.ll
+; RUN: opt -S -mtriple=amdgcn-- -amdgpu-lower-module-lds %t.ll -o -
+; RUN: diff -ub %t.ll %t.second.ll -I ".*ModuleID.*"
+
+; Check AMDGPULowerModuleLDS can run more than once on the same module, and that
+; the second run is a no-op.
+
+@lds = internal unnamed_addr addrspace(3) global i32 undef, align 4
+
+define amdgpu_kernel void @test() {
+entry:
+  store i32 1, ptr addrspace(3) @lds
+  ret void
+}