[AMDGPU][SplitModule] Handle !callees metadata

Author: Pierre-vh

llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp

@@ -482,6 +482,27 @@ void SplitGraph::Node::visitAllDependencies(
   }
 }
 
+/// Checks if \p I has MD_callees and if it does, parse it and put the function
+/// in \p Callees.
+///
+/// \returns true if there was metadata and it was parsed correctly. false if
+/// there was no MD or if it contained unknown entries.
+static bool handleCalleesMD(const Instruction &I,
+                            SmallVector<Function *> &Callees) {
+  auto *MD = I.getMetadata(LLVMContext::MD_callees);
+  if (!MD)
+    return false;
+
+  for (const auto &Op : MD->operands()) {
+    Function *Callee = mdconst::extract_or_null<Function>(Op);
+    if (!Callee)
+      return false;
+    Callees.push_back(Callee);
+  }
+
+  return true;
+}
+
 void SplitGraph::buildGraph(CallGraph &CG) {
   SplitModuleTimer SMT("buildGraph", "graph construction");
   LLVM_DEBUG(
@@ -519,6 +540,8 @@ void SplitGraph::buildGraph(CallGraph &CG) {
                  Fn.printAsOperand(dbgs());
                  dbgs() << " - analyzing function\n");
 
+      SmallVector<Function *> KnownCallees;
+
       bool HasIndirectCall = false;
       for (const auto &Inst : instructions(Fn)) {
         // look at all calls without a direct callee.
@@ -531,6 +554,9 @@ void SplitGraph::buildGraph(CallGraph &CG) {
             continue;
           }
 
+          if (handleCalleesMD(Inst, KnownCallees))
+            continue;
+
           // everything else is handled conservatively.
           HasIndirectCall = true;
           break;
@@ -540,7 +566,8 @@ void SplitGraph::buildGraph(CallGraph &CG) {
       if (HasIndirectCall) {
         LLVM_DEBUG(dbgs() << "    indirect call found\n");
         FnsWithIndirectCalls.push_back(&Fn);
-      }
+      } else if (!KnownCallees.empty())
+        DirectCallees.insert(KnownCallees.begin(), KnownCallees.end());
     }
 
     Node &N = getNode(Cache, Fn);

llvm/test/tools/llvm-split/AMDGPU/kernels-dependency-indirect-callee-md.ll

@@ -0,0 +1,69 @@
+; RUN: sed -s 's/_MD_/, !callees !{ptr @CallCandidate0}/' %s | llvm-split -o %t -j 3 -mtriple amdgcn-amd-amdhsa
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 --implicit-check-not=define %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 --implicit-check-not=define %s
+; RUN: llvm-dis -o - %t2 | FileCheck --check-prefix=CHECK2 --implicit-check-not=define %s
+
+; RUN: sed -s 's/_MD_//g' %s | llvm-split -o %t-nomd -j 3 -mtriple amdgcn-amd-amdhsa
+; RUN: llvm-dis -o - %t-nomd0 | FileCheck --check-prefix=CHECK-NOMD0 --implicit-check-not=define %s
+; RUN: llvm-dis -o - %t-nomd1 | FileCheck --check-prefix=CHECK-NOMD1 --implicit-check-not=define %s
+; RUN: llvm-dis -o - %t-nomd2 | FileCheck --check-prefix=CHECK-NOMD2 --implicit-check-not=define %s
+
+; CHECK0: define internal void @HelperC
+; CHECK0: define amdgpu_kernel void @C
+
+; CHECK1: define hidden void @CallCandidate1
+; CHECK1: define internal void @HelperB
+; CHECK1: define amdgpu_kernel void @B
+
+; CHECK2: define internal void @HelperA
+; CHECK2: define hidden void @CallCandidate0
+; CHECK2: define amdgpu_kernel void @A
+
+; CHECK-NOMD0: define internal void @HelperC
+; CHECK-NOMD0: define amdgpu_kernel void @C
+
+; CHECK-NOMD1: define internal void @HelperB
+; CHECK-NOMD1: define amdgpu_kernel void @B
+
+; CHECK-NOMD2: define internal void @HelperA
+; CHECK-NOMD2: define hidden void @CallCandidate0
+; CHECK-NOMD2: define hidden void @CallCandidate1
+; CHECK-NOMD2: define amdgpu_kernel void @A
+
+@addrthief = global [2 x ptr] [ptr @CallCandidate0, ptr @CallCandidate1]
+
+define internal void @HelperA(ptr %call) {
+  call void %call() _MD_
+  ret void
+}
+
+define internal void @CallCandidate0() {
+  ret void
+}
+
+define internal void @CallCandidate1() {
+  ret void
+}
+
+define internal void @HelperB() {
+  ret void
+}
+
+define internal void @HelperC() {
+  ret void
+}
+
+define amdgpu_kernel void @A(ptr %call) {
+  call void @HelperA(ptr %call)
+  ret void
+}
+
+define amdgpu_kernel void @B() {
+  call void @HelperB()
+  ret void
+}
+
+define amdgpu_kernel void @C() {
+  call void @HelperC()
+  ret void
+}