[SYCL] Duplicate address-taken functions in split device modules. (#6452)

* [SYCL] Duplicate address-taken functions in split device code modules.

We conservatively assume that address-taken function can be used
in any split module, so it needs to be duplicated.

Signed-off-by: Konstantin S Bobrovsky <konstantin.s.bobrovsky@intel.com>

Author: kbobrovs

llvm/test/tools/sycl-post-link/split-with-func-ptrs.ll

@@ -0,0 +1,76 @@
+; This test checks that functions which are neither SYCL external functions nor
+; part of any call graph, but have their address taken, are retained in split
+; modules.
+
+; -- Per-source split
+; RUN: sycl-post-link -split=source -emit-param-info -symbols -emit-exported-symbols -split-esimd -lower-esimd -O2 -spec-const=rt -S %s -o %tA.table
+; RUN: FileCheck %s -input-file=%tA_0.ll --check-prefixes CHECK-A0
+; RUN: FileCheck %s -input-file=%tA_1.ll --check-prefixes CHECK-A1
+; -- No split
+; RUN: sycl-post-link -emit-param-info -symbols -emit-exported-symbols -split-esimd -lower-esimd -O2 -spec-const=rt -S %s -o %tB.table
+; RUN: FileCheck %s -input-file=%tB_0.ll --check-prefixes CHECK-B0
+; -- Per-kernel split
+; RUN: sycl-post-link -split=kernel -emit-param-info -symbols -emit-exported-symbols -split-esimd -lower-esimd -O2 -spec-const=rt -S %s -o %tC.table
+; RUN: FileCheck %s -input-file=%tC_0.ll --check-prefixes CHECK-C0
+; RUN: FileCheck %s -input-file=%tC_1.ll --check-prefixes CHECK-C1
+
+
+; ModuleID = 'in.bc'
+source_filename = "llvm-link"
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+$foo = comdat any
+$bar = comdat any
+
+@"tableX" = weak global [1 x void ()*] [void ()* @foo], align 8
+@"tableY" = weak global [1 x void ()*] [void ()* @bar], align 8
+
+
+; Function Attrs: mustprogress norecurse nounwind
+define linkonce_odr dso_local spir_func void @foo() unnamed_addr #0 comdat align 2 {
+; CHECK-A0: define linkonce_odr dso_local spir_func void @foo
+; CHECK-A1: define linkonce_odr dso_local spir_func void @foo
+; CHECK-B0: define linkonce_odr dso_local spir_func void @foo
+; CHECK-B1: define linkonce_odr dso_local spir_func void @foo
+; CHECK-C0: define linkonce_odr dso_local spir_func void @foo
+; CHECK-C1: define linkonce_odr dso_local spir_func void @foo
+  ret void
+}
+
+; Function Attrs: mustprogress norecurse nounwind
+define linkonce_odr dso_local spir_func void @bar() unnamed_addr #1 comdat align 2 {
+; CHECK-A0: define linkonce_odr dso_local spir_func void @bar
+; CHECK-A1: define linkonce_odr dso_local spir_func void @bar
+; CHECK-B0: define linkonce_odr dso_local spir_func void @bar
+; CHECK-B1: define linkonce_odr dso_local spir_func void @bar
+; CHECK-C0: define linkonce_odr dso_local spir_func void @bar
+; CHECK-C1: define linkonce_odr dso_local spir_func void @bar
+  ret void
+}
+
+define weak_odr dso_local spir_kernel void @Kernel1() #2 {
+  ret void
+}
+
+define weak_odr dso_local spir_kernel void @Kernel2() #3 {
+  ret void
+}
+
+attributes #0 = { mustprogress norecurse nounwind "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "vector_function_ptrs"="tableX()" }
+attributes #1 = { mustprogress norecurse nounwind "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "vector_function_ptrs"="tableY()" }
+attributes #2 = { "sycl-module-id"="module1.cpp" "uniform-work-group-size"="true" }
+attributes #3 = { "sycl-module-id"="module2.cpp" "uniform-work-group-size"="true" }
+
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}
+!2 = !{}
+!3 = !{!"<ID>"}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{i32 7, !"frame-pointer", i32 2}
+!6 = !{!7, !8, i64 8}
+!7 = !{!"_ZTS4Base", !8, i64 8}
+!8 = !{!"int", !9, i64 0}
+!9 = !{!"omnipotent char", !10, i64 0}
+!10 = !{!"Simple C++ TBAA"}

llvm/tools/sycl-post-link/ModuleSplitter.cpp

@@ -107,6 +107,10 @@ bool isSpirvSyclBuiltin(StringRef FName) {
   return FName.startswith("__spirv_") || FName.startswith("__sycl_");
 }
 
+bool isKernel(const Function &F) {
+  return F.getCallingConv() == CallingConv::SPIR_KERNEL;
+}
+
 bool isEntryPoint(const Function &F, bool EmitOnlyKernelsAsEntryPoints) {
   // Skip declarations, if any: they should not be included into a vector of
   // entry points groups or otherwise we will end up with incorrectly generated
@@ -115,7 +119,7 @@ bool isEntryPoint(const Function &F, bool EmitOnlyKernelsAsEntryPoints) {
     return false;
 
   // Kernels are always considered to be entry points
-  if (CallingConv::SPIR_KERNEL == F.getCallingConv())
+  if (isKernel(F))
     return true;
 
   if (!EmitOnlyKernelsAsEntryPoints) {
@@ -285,6 +289,7 @@ class CallGraph {
 private:
   std::unordered_map<const Function *, FunctionSet> Graph;
   SmallPtrSet<const Function *, 1> EmptySet;
+  FunctionSet AddrTakenFunctions;
 
 public:
   CallGraph(const Module &M) {
@@ -297,6 +302,9 @@ class CallGraph {
           }
         }
       }
+      if (F.hasAddressTaken()) {
+        AddrTakenFunctions.insert(&F);
+      }
     }
   }
 
@@ -307,13 +315,28 @@ class CallGraph {
                ? make_range(EmptySet.begin(), EmptySet.end())
                : make_range(It->second.begin(), It->second.end());
   }
+
+  iterator_range<FunctionSet::const_iterator> addrTakenFunctions() const {
+    return make_range(AddrTakenFunctions.begin(), AddrTakenFunctions.end());
+  }
 };
 
 void collectFunctionsToExtract(SetVector<const GlobalValue *> &GVs,
                                const EntryPointGroup &ModuleEntryPoints,
                                const CallGraph &Deps) {
   for (const auto *F : ModuleEntryPoints.Functions)
     GVs.insert(F);
+  // It is conservatively assumed that any address-taken function can be invoked
+  // or otherwise used by any function in any module split from the initial one.
+  // So such functions along with the call graphs they start are always
+  // extracted (and duplicated in each split module).
+  // TODO: try to determine which split modules really use address-taken
+  // functions and only duplicate the functions in such modules. Note that usage
+  // may include e.g. function address comparison w/o actual invocation.
+  for (const auto *F : Deps.addrTakenFunctions()) {
+    if (!isKernel(*F) && (isESIMDFunction(*F) == ModuleEntryPoints.isEsimd()))
+      GVs.insert(F);
+  }
 
   // GVs has SetVector type. This type inserts a value only if it is not yet
   // present there. So, recursion is not expected here.