[AMDGPU] Fix missing lowering of LDS used in global scope.

Reviewed By: rampitec

Differential Revision: https://reviews.llvm.org/D103431

3 files changed, 58 insertions, 33 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index e3287f07..937daea 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -171,7 +171,7 @@ private:
 
     // Find variables to move into new struct instance
     std::vector<GlobalVariable *> FoundLocalVars =
-        AMDGPU::findVariablesToLower(M, UsedList, F);
+        AMDGPU::findVariablesToLower(M, F);
 
     if (FoundLocalVars.empty()) {
       // No variables to rewrite, no changes made.
diff --git a/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.cpp b/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.cpp
index ba3e0de..fd704fa 100644
--- a/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.cpp
+++ b/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.cpp
@@ -60,15 +60,34 @@ void replaceConstantUsesInFunction(ConstantExpr *C, const Function *F) {
   }
 }
 
-bool shouldLowerLDSToStruct(const SmallPtrSetImpl<GlobalValue *> &UsedList,
-                            const GlobalVariable &GV, const Function *F) {
-  // Any LDS variable can be lowered by moving into the created struct
-  // Each variable so lowered is allocated in every kernel, so variables
-  // whose users are all known to be safe to lower without the transform
-  // are left unchanged.
+bool hasUserInstruction(const GlobalValue *GV) {
+  SmallPtrSet<const User *, 8> Visited;
+  SmallVector<const User *, 16> Stack(GV->users());
+
+  while (!Stack.empty()) {
+    const User *U = Stack.pop_back_val();
+
+    if (!Visited.insert(U).second)
+      continue;
+
+    if (isa<Instruction>(U))
+      return true;
+
+    append_range(Stack, U->users());
+  }
+
+  return false;
+}
+
+bool shouldLowerLDSToStruct(const GlobalVariable &GV, const Function *F) {
+  // We are not interested in kernel LDS lowering for module LDS itself.
+  if (F && GV.getName() == "llvm.amdgcn.module.lds")
+    return false;
+
   bool Ret = false;
   SmallPtrSet<const User *, 8> Visited;
   SmallVector<const User *, 16> Stack(GV.users());
+  SmallPtrSet<const GlobalValue *, 8> GlobalUsers;
 
   assert(!F || isKernelCC(F));
 
@@ -76,10 +95,16 @@ bool shouldLowerLDSToStruct(const SmallPtrSetImpl<GlobalValue *> &UsedList,
     const User *V = Stack.pop_back_val();
     Visited.insert(V);
 
-    if (auto *G = dyn_cast<GlobalValue>(V->stripPointerCasts())) {
-      if (UsedList.contains(G)) {
-        continue;
+    if (auto *G = dyn_cast<GlobalValue>(V)) {
+      StringRef GName = G->getName();
+      if (F && GName != "llvm.used" && GName != "llvm.compiler.used") {
+        // For kernel LDS lowering, if G is not a compiler.used list, then we
+        // cannot lower the lds GV since we cannot replace the use of GV within
+        // G.
+        return false;
       }
+      GlobalUsers.insert(G);
+      continue;
     }
 
     if (auto *I = dyn_cast<Instruction>(V)) {
@@ -88,32 +113,32 @@ bool shouldLowerLDSToStruct(const SmallPtrSetImpl<GlobalValue *> &UsedList,
         // Used from this kernel, we want to put it into the structure.
         Ret = true;
       } else if (!F) {
+        // For module LDS lowering, lowering is required if the user instruction
+        // is from non-kernel function.
         Ret |= !isKernelCC(UF);
       }
       continue;
     }
 
-    if (auto *E = dyn_cast<ConstantExpr>(V)) {
-      for (const User *U : E->users()) {
-        if (Visited.insert(U).second) {
-          Stack.push_back(U);
-        }
-      }
-      continue;
-    }
+    // User V should be a constant, recursively visit users of V.
+    assert(isa<Constant>(V) && "Expected a constant.");
+    append_range(Stack, V->users());
+  }
 
-    // Unknown user, conservatively lower the variable.
-    // For module LDS conservatively means place it into the module LDS struct.
-    // For kernel LDS it means lower as a standalone variable.
-    return !F;
+  if (!F && !Ret) {
+    // For module LDS lowering, we have not yet decided if we should lower GV or
+    // not. Explore all global users of GV, and check if atleast one of these
+    // global users appear as an use within an instruction (possibly nested use
+    // via constant expression), if so, then conservately lower LDS.
+    for (auto *G : GlobalUsers)
+      Ret |= hasUserInstruction(G);
   }
 
   return Ret;
 }
 
-std::vector<GlobalVariable *>
-findVariablesToLower(Module &M, const SmallPtrSetImpl<GlobalValue *> &UsedList,
-                     const Function *F) {
+std::vector<GlobalVariable *> findVariablesToLower(Module &M,
+                                                   const Function *F) {
   std::vector<llvm::GlobalVariable *> LocalVars;
   for (auto &GV : M.globals()) {
     if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) {
@@ -137,7 +162,7 @@ findVariablesToLower(Module &M, const SmallPtrSetImpl<GlobalValue *> &UsedList,
       // dropped by the back end if not. This pass skips over it.
       continue;
     }
-    if (!shouldLowerLDSToStruct(UsedList, GV, F)) {
+    if (!shouldLowerLDSToStruct(GV, F)) {
       continue;
     }
     LocalVars.push_back(&GV);
diff --git a/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.h b/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.h
index b5e2cb9..95011ee 100644
--- a/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.h
+++ b/llvm/lib/Target/AMDGPU/Utils/AMDGPULDSUtils.h
@@ -25,19 +25,19 @@ bool isKernelCC(const Function *Func);
 
 Align getAlign(DataLayout const &DL, const GlobalVariable *GV);
 
+/// \returns true if a given global variable \p GV (or its global users) appear
+/// as an use within some instruction (either from kernel or from non-kernel).
+bool hasUserInstruction(const GlobalValue *GV);
+
 /// \returns true if an LDS global requres lowering to a module LDS structure
 /// if \p F is not given. If \p F is given it must be a kernel and function
 /// \returns true if an LDS global is directly used from that kernel and it
 /// is safe to replace its uses with a kernel LDS structure member.
-/// \p UsedList contains a union of llvm.used and llvm.compiler.used variables
-/// which do not count as a use.
-bool shouldLowerLDSToStruct(const SmallPtrSetImpl<GlobalValue *> &UsedList,
-                            const GlobalVariable &GV,
+bool shouldLowerLDSToStruct(const GlobalVariable &GV,
                             const Function *F = nullptr);
 
-std::vector<GlobalVariable *>
-findVariablesToLower(Module &M, const SmallPtrSetImpl<GlobalValue *> &UsedList,
-                     const Function *F = nullptr);
+std::vector<GlobalVariable *> findVariablesToLower(Module &M,
+                                                   const Function *F = nullptr);
 
 SmallPtrSet<GlobalValue *, 32> getUsedList(Module &M);