Remove internal functions made unreachable by predicate expansion.

Author: AlexVlx

llvm/lib/Target/AMDGPU/AMDGPUExpandFeaturePredicates.cpp

@@ -127,6 +127,30 @@ handlePredicate(const GCNSubtarget &ST, FunctionAnalysisManager &FAM,
 }
 } // Unnamed namespace.
 
+static inline SmallVector<Function *> collectUsedFunctions(Module &M) {
+  SmallVector<Function *> Ret;
+  for (auto &&F : M) {
+    if (F.isIntrinsic() || F.isDeclaration())
+      continue;
+    if (!F.hasInternalLinkage() && !F.hasPrivateLinkage())
+      continue;
+    if (F.hasNUndroppableUsesOrMore(1))
+      Ret.push_back(&F);
+  }
+
+  return Ret;
+}
+
+template<typename Container0, typename Container1, typename Container2>
+static inline void removeUnreachable(const Container0 &Predicates,
+                                     const Container1 &PredicatedFns,
+                                     const Container2 &UnreachableFns) {
+  for_each(Predicates, [](auto &&P) { P->eraseFromParent(); });
+  for_each(PredicatedFns, [](auto &&F) { removeUnreachableBlocks(*F); });
+  for_each(UnreachableFns,
+           [](auto &&F) { if (F->getNumUses() == 0) F->eraseFromParent(); });
+}
+
 PreservedAnalyses
 AMDGPUExpandFeaturePredicatesPass::run(Module &M, ModuleAnalysisManager &MAM) {
   if (M.empty())
@@ -148,6 +172,7 @@ AMDGPUExpandFeaturePredicatesPass::run(Module &M, ModuleAnalysisManager &MAM) {
 
   auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   SmallPtrSet<Function *, 32> Predicated;
+  SmallVector<Function *> MaybeUnreachable = collectUsedFunctions(M);
   auto Ret = PreservedAnalyses::all();
   for (auto &&P : Predicates) {
     auto R = handlePredicate(ST, FAM, Predicated, P);
@@ -158,10 +183,7 @@ AMDGPUExpandFeaturePredicatesPass::run(Module &M, ModuleAnalysisManager &MAM) {
     Ret.intersect(R.first);
   }
 
-  for (auto &&P : Predicates)
-    P->eraseFromParent();
-  for (auto &&F : Predicated)
-    removeUnreachableBlocks(*F);
+  removeUnreachable(Predicates, Predicated, MaybeUnreachable);
 
   return Ret;
 }

llvm/test/CodeGen/AMDGPU/amdgpu-expand-feature-predicates-remove-unreachable-functions.ll

@@ -0,0 +1,104 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --scrub-attributes --version 5
+; REQUIRES: amdgpu-registered-target
+
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 -passes='amdgpu-expand-feature-predicates' %s -o - | FileCheck --check-prefix=GFX906 %s
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1201 -passes='amdgpu-expand-feature-predicates' %s -o - | FileCheck --check-prefix=GFX1201 %s
+
+@llvm.amdgcn.is.gfx906 = external addrspace(1) externally_initialized constant i1
+@llvm.amdgcn.is.gfx1201 = external addrspace(1) externally_initialized constant i1
+
+define external void @extern_linkage() {
+; GFX906-LABEL: define void @extern_linkage(
+; GFX906-SAME: ) #[[ATTR0:[0-9]+]] {
+; GFX906-NEXT:  [[ENTRY:.*:]]
+; GFX906-NEXT:    ret void
+;
+; GFX1201-LABEL: define void @extern_linkage(
+; GFX1201-SAME: ) #[[ATTR0:[0-9]+]] {
+; GFX1201-NEXT:  [[ENTRY:.*:]]
+; GFX1201-NEXT:    ret void
+;
+entry:
+  ret void
+}
+
+define private void @non_predicated_uses() {
+; GFX906-LABEL: define private void @non_predicated_uses(
+; GFX906-SAME: ) #[[ATTR0]] {
+; GFX906-NEXT:  [[ENTRY:.*:]]
+; GFX906-NEXT:    ret void
+;
+; GFX1201-LABEL: define private void @non_predicated_uses(
+; GFX1201-SAME: ) #[[ATTR0]] {
+; GFX1201-NEXT:  [[ENTRY:.*:]]
+; GFX1201-NEXT:    ret void
+;
+entry:
+  ret void
+}
+
+define internal void @remove_on_906() {
+; GFX1201-LABEL: define internal void @remove_on_906(
+; GFX1201-SAME: ) #[[ATTR0]] {
+; GFX1201-NEXT:  [[ENTRY:.*:]]
+; GFX1201-NEXT:    ret void
+;
+entry:
+  ret void
+}
+
+define internal void @remove_on_1201() {
+; GFX906-LABEL: define internal void @remove_on_1201(
+; GFX906-SAME: ) #[[ATTR0]] {
+; GFX906-NEXT:  [[ENTRY:.*:]]
+; GFX906-NEXT:    ret void
+;
+entry:
+  ret void
+}
+
+define void @foo() {
+; GFX906-LABEL: define void @foo(
+; GFX906-SAME: ) #[[ATTR0]] {
+; GFX906-NEXT:  [[ENTRY:.*:]]
+; GFX906-NEXT:    call void @non_predicated_uses()
+; GFX906-NEXT:    br label %[[NOT_GFX1201:.*]]
+; GFX906:       [[NOT_GFX1201]]:
+; GFX906-NEXT:    br label %[[GFX906:.*]]
+; GFX906:       [[GFX906]]:
+; GFX906-NEXT:    call void @remove_on_1201()
+; GFX906-NEXT:    br label %[[END:.*]]
+; GFX906:       [[END]]:
+; GFX906-NEXT:    ret void
+;
+; GFX1201-LABEL: define void @foo(
+; GFX1201-SAME: ) #[[ATTR0]] {
+; GFX1201-NEXT:  [[ENTRY:.*:]]
+; GFX1201-NEXT:    call void @non_predicated_uses()
+; GFX1201-NEXT:    br label %[[GFX1201:.*]]
+; GFX1201:       [[GFX1201]]:
+; GFX1201-NEXT:    call void @remove_on_906()
+; GFX1201-NEXT:    br label %[[END:.*]]
+; GFX1201:       [[END]]:
+; GFX1201-NEXT:    ret void
+;
+entry:
+  call void @non_predicated_uses()
+  %0 = load i1, ptr addrspace(1) @llvm.amdgcn.is.gfx1201, align 1
+  br i1 %0, label %gfx1201, label %not.gfx1201
+
+gfx1201:
+  call void @remove_on_906()
+  br label %end
+
+not.gfx1201:
+  %1 = load i1, ptr addrspace(1) @llvm.amdgcn.is.gfx906, align 1
+  br i1 %1, label %gfx906, label %end
+
+gfx906:
+  call void @remove_on_1201()
+  br label %end
+
+end:
+  ret void
+}