[clang] [llvm] [clang-sycl-linker] Add per-kernel device code splitting (PR #195362)

[Yury Plyakhin via cfe-commits]

================
@@ -461,6 +468,107 @@ static Error runAOTCompile(StringRef InputFile, StringRef OutputFile,
   return createStringError(inconvertibleErrorCode(), "Unsupported arch");
 }
 
+/// SYCL device code module split mode.
+enum IRSplitMode {
+  SPLIT_PER_KERNEL, // one module per kernel
+  SPLIT_NONE        // no splitting
+};
+
+/// Parses the value of \p -sycl-module-split-mode.
+static std::optional<IRSplitMode> convertStringToSplitMode(StringRef S) {
+  return StringSwitch<std::optional<IRSplitMode>>(S)
+      .Case("kernel", SPLIT_PER_KERNEL)
+      .Case("none", SPLIT_NONE)
+      .Default(std::nullopt);
+}
+
+/// Result of splitting a device module: the bitcode file path and the
+/// serialized symbol table for each device image.
+struct SplitModule {
+  SmallString<256> ModuleFilePath;
+  SmallString<0> Symbols;
+};
+
+static bool isEntryPoint(const Function &F) {
+  return !F.isDeclaration() && F.hasKernelCallingConv();
+}
+
+/// Collect kernel names from \p M and serialize them into a symbol table.
+static SmallString<0> collectSymbols(const Module &M) {
+  SmallVector<StringRef> KernelNames;
+  for (const Function &F : M)
+    if (isEntryPoint(F))
+      KernelNames.push_back(F.getName());
+  SmallString<0> SymbolData;
+  llvm::offloading::sycl::writeSymbolTable(KernelNames, SymbolData);
+  return SymbolData;
+}
+
+/// Splits the fully linked device \p M into one bitcode file per device image
+/// according to \p Mode and returns the list of split images with their symbol
+/// tables.
+///
+/// For SPLIT_NONE, \p LinkedBitcodeFile is returned as-is.
+/// For SPLIT_PER_KERNEL, the module is split into parts such that each part
+/// contains exactly one kernel entry point and its transitive dependencies;
+/// each part is written to a fresh temporary bitcode file.
+static Expected<SmallVector<SplitModule, 0>>
+splitDeviceCode(std::unique_ptr<Module> M, StringRef LinkedBitcodeFile,
+                IRSplitMode Mode, const ArgList &Args) {
+  SmallVector<SplitModule, 0> SplitModules;
+
+  if (Mode == SPLIT_NONE) {
+    SplitModules.push_back(
+        {SmallString<256>(LinkedBitcodeFile), collectSymbols(*M)});
+    return SplitModules;
+  }
+
+  assert(Mode == SPLIT_PER_KERNEL);
+
+  // splitModuleTransitiveFromEntryPoints asserts that at least one entry point
+  // was categorized. If the linked module contains no kernel definitions at
+  // all, there is nothing to split; fall back to shipping the linked module
+  // as a single image.
+  bool HasKernel = llvm::any_of(M->functions(), isEntryPoint);
+  if (!HasKernel) {
+    SplitModules.push_back(
+        {SmallString<256>(LinkedBitcodeFile), collectSymbols(*M)});
+    return SplitModules;
+  }
+
+  // Categorize each kernel function into its own group. Non-kernels and
+  // declarations return std::nullopt so they are pulled into whichever split
+  // transitively needs them.
+  int NextCategory = 0;
+  auto EntryPointCategorizer =
+      [&NextCategory](const Function &F) -> std::optional<int> {
+    if (!isEntryPoint(F))
+      return std::nullopt;
+    return NextCategory++;
+  };
+
+  if (Error Err = splitModuleTransitiveFromEntryPoints(
+          std::move(M), EntryPointCategorizer,
+          [&](std::unique_ptr<Module> Part) -> Error {
----------------
YuriPlyakhin wrote:

https://github.com/llvm/llvm-project/pull/195362/changes/33e58bf71528c90601dd9e36dec82a9534513bd1

https://github.com/llvm/llvm-project/pull/195362