[llvm] [AMDGPU] Let LowerModuleLDS run twice on the same module (PR #81729)

[Pierre van Houtryve via llvm-commits]

https://github.com/Pierre-vh created https://github.com/llvm/llvm-project/pull/81729

If all variables in the module are absolute, this means we're running the pass again on an already lowered module, and that works.
If none of them are absolute, lowering can proceed as usual.
Only diagnose cases where we have a mix of absolute/non-absolute GVs, which means we added LDS GVs after lowering, which is broken.

See #81491
Split from #75333

>From 58359965fbd0bf2bf1f1d6001e4643b062a7cba4 Mon Sep 17 00:00:00 2001
From: pvanhout <pierre.vanhoutryve at amd.com>
Date: Wed, 14 Feb 2024 13:08:32 +0100
Subject: [PATCH] [AMDGPU] Let LowerModuleLDS run twice on the same module

See #81491
Split from #75333
---
 .../AMDGPU/AMDGPULowerModuleLDSPass.cpp       | 26 +++++++++++++------
 ...=> lds-reject-mixed-absolute-addresses.ll} |  4 +--
 llvm/test/CodeGen/AMDGPU/lds-run-twice.ll     | 14 ++++++++++
 3 files changed, 34 insertions(+), 10 deletions(-)
 rename llvm/test/CodeGen/AMDGPU/{lds-reject-absolute-addresses.ll => lds-reject-mixed-absolute-addresses.ll} (81%)
 create mode 100644 llvm/test/CodeGen/AMDGPU/lds-run-twice.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index 5762f1906a16d3..38475be1ec4726 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -327,6 +327,8 @@ class AMDGPULowerModuleLDS {
     return convertUsersOfConstantsToInstructions(LDSGlobals);
   }
 
+  std::optional<bool> HasAbsoluteGVs;
+
 public:
   AMDGPULowerModuleLDS(const AMDGPUTargetMachine &TM_) : TM(TM_) {}
 
@@ -334,9 +336,9 @@ class AMDGPULowerModuleLDS {
 
   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;
diff --git a/llvm/test/CodeGen/AMDGPU/lds-reject-absolute-addresses.ll b/llvm/test/CodeGen/AMDGPU/lds-reject-mixed-absolute-addresses.ll
similarity index 81%
rename from llvm/test/CodeGen/AMDGPU/lds-reject-absolute-addresses.ll
rename to llvm/test/CodeGen/AMDGPU/lds-reject-mixed-absolute-addresses.ll
index 659cdb55ded2fe..b512a43aa10222 100644
--- a/llvm/test/CodeGen/AMDGPU/lds-reject-absolute-addresses.ll
+++ b/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
+ at 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}
-
diff --git a/llvm/test/CodeGen/AMDGPU/lds-run-twice.ll b/llvm/test/CodeGen/AMDGPU/lds-run-twice.ll
new file mode 100644
index 00000000000000..63aa426c9b9e6e
--- /dev/null
+++ b/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.
+
+ at 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
+}