diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp b/llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp
--- a/llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp
@@ -29,6 +29,7 @@
 #include "AMDGPU.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "Utils/AMDGPUBaseInfo.h"
+#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
@@ -157,12 +158,12 @@
 
   // Remove attributes from F.
   // This is used in presence of address taken functions.
-  bool removeAttributes(Function *F);
+  bool removeAttributes(Function &F);
 
   // Handle call graph rooted at address taken functions.
   // This function will erase all attributes present
   // on all functions called from address taken functions transitively.
-  bool handleAddressTakenFunctions(CallGraph *CG);
+  bool handleAddressTakenFunctions();
 };
 
 // Allows to propagate attributes early, but no clonning is allowed as it must
@@ -216,46 +217,36 @@
                 "Late propagate attributes from kernels to functions",
                 false, false)
 
-bool AMDGPUPropagateAttributes::removeAttributes(Function *F) {
+bool AMDGPUPropagateAttributes::removeAttributes(Function &F) {
   bool Changed = false;
-  if (!F)
-    return Changed;
-  LLVM_DEBUG(dbgs() << "Removing attributes from " << F->getName() << '\n');
-  for (unsigned I = 0; I < NumAttr; ++I) {
-    if (F->hasFnAttribute(AttributeNames[I])) {
-      F->removeFnAttr(AttributeNames[I]);
+  LLVM_DEBUG(dbgs() << "Removing attributes from " << F.getName() << '\n');
+  for (const char *AttrName : AttributeNames) {
+    if (F.hasFnAttribute(AttrName)) {
+      F.removeFnAttr(AttrName);
       Changed = true;
     }
   }
   return Changed;
 }
 
-bool AMDGPUPropagateAttributes::handleAddressTakenFunctions(CallGraph *CG) {
+bool AMDGPUPropagateAttributes::handleAddressTakenFunctions() {
   assert(ModuleCG && "Call graph not present");
 
   bool Changed = false;
-  SmallSet<CallGraphNode *, 32> Visited;
+  df_iterator_default_set<CallGraphNode *> Reachable;
 
+  // Find all functions reachable from address-taken functions.
   for (Function *F : AddressTakenFunctions) {
-    CallGraphNode *CGN = (*CG)[F];
-    if (!Visited.count(CGN)) {
-      Changed |= removeAttributes(F);
-      Visited.insert(CGN);
-    }
+    for (auto *I : depth_first_ext((*ModuleCG)[F], Reachable))
+      (void)I;
+  }
 
-    std::queue<CallGraphNode *> SubGraph;
-    SubGraph.push(CGN);
-    while (!SubGraph.empty()) {
-      CallGraphNode *CGN = SubGraph.front();
-      SubGraph.pop();
-      if (!Visited.count(CGN)) {
-        Changed |= removeAttributes(CGN->getFunction());
-        Visited.insert(CGN);
-      }
-      for (auto N : *CGN)
-        SubGraph.push(N.second);
-    }
+  // Remove attributes from all reachable functions.
+  for (CallGraphNode *CGN : Reachable) {
+    if (Function *F = CGN->getFunction())
+      Changed |= removeAttributes(*F);
   }
+
   return Changed;
 }
 
@@ -378,7 +369,7 @@
   // calls separate to handle them gracefully.
   // The core traversal need not be affected by this.
   if (AllowClone)
-    Changed |= handleAddressTakenFunctions(ModuleCG);
+    Changed |= handleAddressTakenFunctions();
 
   return Changed;
 }