diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index 543e423878877c..2617f0fcf0b2e4 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -1034,21 +1034,18 @@ class AMDGPULowerModuleLDS {
   void removeNoLdsKernelIdFromReachable(CallGraph &CG, Function *KernelRoot) {
     KernelRoot->removeFnAttr("amdgpu-no-lds-kernel-id");
 
-    SmallVector<Function *> Tmp({CG[KernelRoot]->getFunction()});
-    if (!Tmp.back())
-      return;
-
+    SmallVector<Function *> WorkList({CG[KernelRoot]->getFunction()});
     SmallPtrSet<Function *, 8> Visited;
     bool SeenUnknownCall = false;
 
-    do {
-      Function *F = Tmp.pop_back_val();
+    while (!WorkList.empty()) {
+      Function *F = WorkList.pop_back_val();
 
-      for (auto &N : *CG[F]) {
-        if (!N.second)
+      for (auto &CallRecord : *CG[F]) {
+        if (!CallRecord.second)
           continue;
 
-        Function *Callee = N.second->getFunction();
+        Function *Callee = CallRecord.second->getFunction();
         if (!Callee) {
           if (!SeenUnknownCall) {
             SeenUnknownCall = true;
@@ -1056,21 +1053,21 @@ class AMDGPULowerModuleLDS {
             // If we see any indirect calls, assume nothing about potential
             // targets.
             // TODO: This could be refined to possible LDS global users.
-            for (auto &N : *CG.getExternalCallingNode()) {
-              Function *PotentialCallee = N.second->getFunction();
+            for (auto &ExternalCallRecord : *CG.getExternalCallingNode()) {
+              Function *PotentialCallee =
+                  ExternalCallRecord.second->getFunction();
+              assert(PotentialCallee);
               if (!isKernelLDS(PotentialCallee))
                 PotentialCallee->removeFnAttr("amdgpu-no-lds-kernel-id");
             }
-
-            continue;
           }
+        } else {
+          Callee->removeFnAttr("amdgpu-no-lds-kernel-id");
+          if (Visited.insert(Callee).second)
+            WorkList.push_back(Callee);
         }
-
-        Callee->removeFnAttr("amdgpu-no-lds-kernel-id");
-        if (Visited.insert(Callee).second)
-          Tmp.push_back(Callee);
       }
-    } while (!Tmp.empty());
+    }
   }
 
   DenseMap<Function *, GlobalVariable *> lowerDynamicLDSVariables(