[llvm] 9981afd - [amdgpu][nfc] Extract kernel annotation from processUsedLDS

[Jon Chesterfield via llvm-commits]

Author: Jon Chesterfield
Date: 2022-07-27T01:38:41+01:00
New Revision: 9981afdd42fad55f9a1b61bd2f38d4db5fe53250

URL: https://github.com/llvm/llvm-project/commit/9981afdd42fad55f9a1b61bd2f38d4db5fe53250
DIFF: https://github.com/llvm/llvm-project/commit/9981afdd42fad55f9a1b61bd2f38d4db5fe53250.diff

LOG: [amdgpu][nfc] Extract kernel annotation from processUsedLDS

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index 10d7dab947bc..96a34caa8eb0 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -149,11 +149,42 @@ class AMDGPULowerModuleLDS : public ModulePass {
   }
 
   bool runOnModule(Module &M) override {
+    LLVMContext &Ctx = M.getContext();
     CallGraph CG = CallGraph(M);
     bool Changed = superAlignLDSGlobals(M);
+
     std::vector<GlobalVariable *> ModuleScopeVariables =
         AMDGPU::findVariablesToLower(M, nullptr);
-    Changed |= processUsedLDS(CG, M, ModuleScopeVariables);
+    if (!ModuleScopeVariables.empty()) {
+      GlobalVariable *SGV =
+          processUsedLDS(CG, M, ModuleScopeVariables, nullptr);
+
+      // This ensures the variable is allocated when called functions access it.
+      // It also lets other passes, specifically PromoteAlloca, accurately
+      // calculate how much LDS will be used by the kernel after lowering.
+
+      IRBuilder<> Builder(Ctx);
+      for (Function &Func : M.functions()) {
+        if (!Func.isDeclaration() && AMDGPU::isKernelCC(&Func)) {
+          const CallGraphNode *N = CG[&Func];
+          const bool CalleesRequireModuleLDS = N->size() > 0;
+
+          if (CalleesRequireModuleLDS) {
+            // If a function this kernel might call requires module LDS,
+            // annotate the kernel to let later passes know it will allocate
+            // this structure, even if not apparent from the IR.
+            markUsedByKernel(Builder, &Func, SGV);
+          } else {
+            // However if we are certain this kernel cannot call a function that
+            // requires module LDS, annotate the kernel so the backend can elide
+            // the allocation without repeating callgraph walks.
+            Func.addFnAttr("amdgpu-elide-module-lds");
+          }
+        }
+      }
+
+      Changed = true;
+    }
 
     for (Function &F : M.functions()) {
       if (F.isDeclaration())
@@ -162,9 +193,13 @@ class AMDGPULowerModuleLDS : public ModulePass {
       // Only lower compute kernels' LDS.
       if (!AMDGPU::isKernel(F.getCallingConv()))
         continue;
+
       std::vector<GlobalVariable *> KernelUsedVariables =
           AMDGPU::findVariablesToLower(M, &F);
-      Changed |= processUsedLDS(CG, M, KernelUsedVariables, &F);
+      if (!KernelUsedVariables.empty()) {
+        processUsedLDS(CG, M, KernelUsedVariables, &F);
+        Changed = true;
+      }
     }
 
     return Changed;
@@ -306,17 +341,13 @@ class AMDGPULowerModuleLDS : public ModulePass {
     return {SGV, std::move(Map)};
   }
 
-  bool processUsedLDS(CallGraph const &CG, Module &M,
-                      std::vector<GlobalVariable *> const &LDSVarsToTransform,
-                      Function *F = nullptr) {
+  GlobalVariable *
+  processUsedLDS(CallGraph const &CG, Module &M,
+                 std::vector<GlobalVariable *> const &LDSVarsToTransform,
+                 Function *F) {
     LLVMContext &Ctx = M.getContext();
     const DataLayout &DL = M.getDataLayout();
 
-    if (LDSVarsToTransform.empty()) {
-      // No variables to rewrite, no changes made.
-      return false;
-    }
-
     std::string VarName(
         F ? (Twine("llvm.amdgcn.kernel.") + F->getName() + ".lds").str()
           : "llvm.amdgcn.module.lds");
@@ -396,31 +427,7 @@ class AMDGPULowerModuleLDS : public ModulePass {
       refineUsesAlignmentAndAA(GEP, A, DL, AliasScope, NoAlias);
     }
 
-    // This ensures the variable is allocated when called functions access it.
-    // It also lets other passes, specifically PromoteAlloca, accurately
-    // calculate how much LDS will be used by the kernel after lowering.
-    if (!F) {
-      IRBuilder<> Builder(Ctx);
-      for (Function &Func : M.functions()) {
-        if (!Func.isDeclaration() && AMDGPU::isKernelCC(&Func)) {
-          const CallGraphNode *N = CG[&Func];
-          const bool CalleesRequireModuleLDS = N->size() > 0;
-
-          if (CalleesRequireModuleLDS) {
-            // If a function this kernel might call requires module LDS,
-            // annotate the kernel to let later passes know it will allocate
-            // this structure, even if not apparent from the IR.
-            markUsedByKernel(Builder, &Func, SGV);
-          } else {
-            // However if we are certain this kernel cannot call a function that
-            // requires module LDS, annotate the kernel so the backend can elide
-            // the allocation without repeating callgraph walks.
-            Func.addFnAttr("amdgpu-elide-module-lds");
-          }
-        }
-      }
-    }
-    return true;
+    return SGV;
   }
 
   void refineUsesAlignmentAndAA(Value *Ptr, Align A, const DataLayout &DL,