[llvm] ce51c5d - AMDGPU: Fix crashing on kernel declarations when lowering LDS

[Matt Arsenault via llvm-commits]

Author: Matt Arsenault
Date: 2021-08-26T19:01:10-04:00
New Revision: ce51c5d4a935cb1a4b99e5c95e0c4c5724b64375

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

LOG: AMDGPU: Fix crashing on kernel declarations when lowering LDS

This was trying to insert the used marker into a declaration.

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
    llvm/test/CodeGen/AMDGPU/lower-module-lds.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index 26e2b5ff4d4bd..02bb3c43c9e73 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -163,6 +163,9 @@ class AMDGPULowerModuleLDS : public ModulePass {
     bool Changed = processUsedLDS(M);
 
     for (Function &F : M.functions()) {
+      if (F.isDeclaration())
+        continue;
+
       // Only lower compute kernels' LDS.
       if (!AMDGPU::isKernel(F.getCallingConv()))
         continue;
@@ -347,11 +350,13 @@ class AMDGPULowerModuleLDS : public ModulePass {
     if (!F) {
       IRBuilder<> Builder(Ctx);
       SmallPtrSet<Function *, 32> Kernels;
-      for (auto &I : M.functions()) {
-        Function *Func = &I;
-        if (AMDGPU::isKernelCC(Func) && !Kernels.contains(Func)) {
-          markUsedByKernel(Builder, Func, SGV);
-          Kernels.insert(Func);
+      for (Function &Func : M.functions()) {
+        if (Func.isDeclaration())
+          continue;
+
+        if (AMDGPU::isKernelCC(&Func) && !Kernels.contains(&Func)) {
+          markUsedByKernel(Builder, &Func, SGV);
+          Kernels.insert(&Func);
         }
       }
     }

diff  --git a/llvm/test/CodeGen/AMDGPU/lower-module-lds.ll b/llvm/test/CodeGen/AMDGPU/lower-module-lds.ll
index e92d8624a4bd2..a620eb88a602c 100644
--- a/llvm/test/CodeGen/AMDGPU/lower-module-lds.ll
+++ b/llvm/test/CodeGen/AMDGPU/lower-module-lds.ll
@@ -54,3 +54,7 @@ define amdgpu_kernel void @kern_call() {
 define spir_kernel void @kern_empty() {
   ret void
 }
+
+; Make sure we don't crash trying to insert code into a kernel
+; declaration.
+declare amdgpu_kernel void @kernel_declaration()