[AMDGPU] Run LowerLDS at the end of the fullLTO pipeline

This change allows us to use `--lto-partitions` in some cases (not guaranteed it works perfectly), as LDS is lowered before the module is split for parallel codegen.

LowerrLDS doesn't support being ran twice because
it'll think the lowered LDS GVs are "absolute addresses" LDS which aren't supported, so I just added a module flag to detect multiple runs.

Author: Pierre-vh

llvm/lib/Target/AMDGPU/AMDGPU.h

@@ -130,7 +130,10 @@ extern char &AMDGPULowerModuleLDSLegacyPassID;
 
 struct AMDGPULowerModuleLDSPass : PassInfoMixin<AMDGPULowerModuleLDSPass> {
   const AMDGPUTargetMachine &TM;
-  AMDGPULowerModuleLDSPass(const AMDGPUTargetMachine &TM_) : TM(TM_) {}
+  bool IsEarlyRun;
+  AMDGPULowerModuleLDSPass(const AMDGPUTargetMachine &TM_,
+                           bool IsEarlyRun = false)
+      : TM(TM_), IsEarlyRun(IsEarlyRun) {}
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

@@ -215,6 +215,12 @@ using namespace llvm;
 
 namespace {
 
+cl::opt<bool>
+    ForceAddModuleFlag("amdgpu-lower-module-lds-force-add-moduleflag",
+                       cl::desc("Always add the module flag that prevents "
+                                "multiple runs of LowerModuleLDS."),
+                       cl::init(false), cl::ReallyHidden);
+
 cl::opt<bool> SuperAlignLDSGlobals(
     "amdgpu-super-align-lds-globals",
     cl::desc("Increase alignment of LDS if it is not on align boundary"),
@@ -254,6 +260,7 @@ template <typename T> std::vector<T> sortByName(std::vector<T> &&V) {
 
 class AMDGPULowerModuleLDS {
   const AMDGPUTargetMachine &TM;
+  bool IsEarlyRun;
 
   static void
   removeLocalVarsFromUsedLists(Module &M,
@@ -328,7 +335,8 @@ class AMDGPULowerModuleLDS {
   }
 
 public:
-  AMDGPULowerModuleLDS(const AMDGPUTargetMachine &TM_) : TM(TM_) {}
+  AMDGPULowerModuleLDS(const AMDGPUTargetMachine &TM_, bool IsEarlyRun = false)
+      : TM(TM_), IsEarlyRun(IsEarlyRun) {}
 
   using FunctionVariableMap = DenseMap<Function *, DenseSet<GlobalVariable *>>;
 
@@ -1088,6 +1096,15 @@ class AMDGPULowerModuleLDS {
   }
 
   bool runOnModule(Module &M) {
+    // This pass may run twice in a full LTO pipeline.
+    //
+    // If we ran it early, we'll have added metadata to skip next runs.
+    if (M.getModuleFlag("amdgcn.lowered_module_lds"))
+      return false;
+    if (IsEarlyRun || ForceAddModuleFlag)
+      M.addModuleFlag(Module::ModFlagBehavior::Warning,
+                      "amdgcn.lowered_module_lds", 1);
+
     CallGraph CG = CallGraph(M);
     bool Changed = superAlignLDSGlobals(M);
 
@@ -1574,6 +1591,7 @@ llvm::createAMDGPULowerModuleLDSLegacyPass(const AMDGPUTargetMachine *TM) {
 
 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M,
                                                 ModuleAnalysisManager &) {
-  return AMDGPULowerModuleLDS(TM).runOnModule(M) ? PreservedAnalyses::none()
-                                                 : PreservedAnalyses::all();
+  return AMDGPULowerModuleLDS(TM, IsEarlyRun).runOnModule(M)
+             ? PreservedAnalyses::none()
+             : PreservedAnalyses::all();
 }

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -770,6 +770,15 @@ void AMDGPUTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
 
         PM.addPass(createCGSCCToFunctionPassAdaptor(std::move(FPM)));
       });
+
+  PB.registerFullLinkTimeOptimizationLastEPCallback(
+      [this](ModulePassManager &PM, OptimizationLevel Level) {
+        // We want to support the -lto-partitions=N option as "best effort".
+        // For that, we need to lower LDS earlier in the pipeline before the
+        // module is partitioned for codegen.
+        if (EnableLowerModuleLDS)
+          PM.addPass(AMDGPULowerModuleLDSPass(*this, /*IsEarlyRun*/ true));
+      });
 }
 
 int64_t AMDGPUTargetMachine::getNullPointerValue(unsigned AddrSpace) {

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

@@ -0,0 +1,39 @@
+; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-lower-module-lds-strategy=module %s -o %t.ll
+; RUN: not --crash opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-lower-module-lds-strategy=module %t.ll -o - 2>&1 | FileCheck %s --check-prefix=ERR
+
+; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-lower-module-lds-strategy=module --amdgpu-lower-module-lds-force-add-moduleflag=1 %s -o %t.ll
+; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-lower-module-lds-strategy=module %t.ll -o - | FileCheck %s
+
+; Check re-run of LowerModuleLDS don't crash when the module flag is used.
+;
+; We first check this test still crashes when ran twice. If it no longer crashes at some point
+; we should update it to ensure the flag still does its job.
+;
+; This test jus has the bare minimum checks to see if the pass ran.
+
+; ERR: LLVM ERROR: LDS variables with absolute addresses are unimplemented.
+
+; CHECK: %llvm.amdgcn.module.lds.t = type { float, [4 x i8], i32 }
+; CHECK: @llvm.amdgcn.module.lds = internal addrspace(3) global %llvm.amdgcn.module.lds.t poison, align 8
+
+; CHECK: attributes #0 = { "amdgpu-lds-size"="12" }
+
+@var0 = addrspace(3) global float poison, align 8
+@var1 = addrspace(3) global i32 poison, align 8
+@ptr = addrspace(1) global ptr addrspace(3) @var1, align 4
+@with_init = addrspace(3) global i64 0
+
+define void @func() {
+  %dec = atomicrmw fsub ptr addrspace(3) @var0, float 1.0 monotonic
+  %val0 = load i32, ptr addrspace(3) @var1, align 4
+  %val1 = add i32 %val0, 4
+  store i32 %val1, ptr addrspace(3) @var1, align 4
+  %unused0 = atomicrmw add ptr addrspace(3) @with_init, i64 1 monotonic
+  ret void
+}
+
+define amdgpu_kernel void @kern_call() {
+  call void @func()
+  %dec = atomicrmw fsub ptr addrspace(3) @var0, float 2.0 monotonic
+  ret void
+}