[offload][SYCL] Add Module splitting by categories. (#131347)

This patch adds Module splitting by categories. The splitting algorithm
is the necessary step in the SYCL compilation pipeline. Also it could be
reused for other heterogenous targets.

The previous attempt was at #119713. In this patch there is no
dependency in `TransformUtils` on "IPO" and on "Printing Passes". In
this patch a module splitting is self-contained and it doesn't introduce
linking issues.

llvm/include/llvm/Transforms/Utils/SplitModuleByCategory.h

@@ -0,0 +1,64 @@
+//===-------- SplitModuleByCategory.h - module split ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// Functionality to split a module by categories.
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORM_UTILS_SPLIT_MODULE_BY_CATEGORY_H
+#define LLVM_TRANSFORM_UTILS_SPLIT_MODULE_BY_CATEGORY_H
+
+#include "llvm/ADT/STLFunctionalExtras.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+
+namespace llvm {
+
+class Module;
+class Function;
+
+/// Splits the given module \p M into parts. Each output part is passed to
+/// \p Callback for further possible processing. Each part corresponds to a
+/// subset of the module that is transitively reachable from some entry point
+/// group. Each entry point group is defined by \p EntryPointCategorizer (EPC)
+/// as follows: 1) If the function is not an entry point, then the Categorizer
+/// returns std::nullopt. Therefore, the function doesn't belong to any group.
+/// However, the function and global objects can still be associated with some
+/// output parts if they are transitively used from some entry points. 2) If the
+/// function belongs to an entry point group, then EPC returns an integer which
+/// is an identifier of the group. If two entry points belong to one group, then
+/// EPC returns the same identifier for both of them.
+///
+/// Let A and B be global objects in the module. The transitive dependency
+/// relation is defined such that: If global object A is used by global object B
+/// in any way (e.g., store, bitcast, phi node, call), then "A" -> "B".
+/// Transitivity is defined such that: If "A" -> "B" and "B" -> "C", then "A" ->
+/// "C". Examples of dependencies:
+/// - Function FA calls function FB
+/// - Function FA uses global variable GA
+/// - Global variable GA references (is initialized with) function FB
+/// - Function FA stores the address of function FB somewhere
+///
+/// The following cases are treated as dependencies between global objects:
+/// 1. Global object A is used by global object B in any way (store,
+///    bitcast, phi node, call, etc.): an "A" -> "B" edge will be added to the
+///    graph;
+/// 2. Function A performs an indirect call of a function with signature S, and
+///    there is a function B with signature S. An "A" -> "B" edge will be added
+///    to the graph;
+///
+/// FIXME: For now, the algorithm assumes no recursion in the input Module. This
+/// will be addressed in the near future.
+void splitModuleTransitiveFromEntryPoints(
+    std::unique_ptr<Module> M,
+    function_ref<std::optional<int>(const Function &F)> EntryPointCategorizer,
+    function_ref<void(std::unique_ptr<Module> Part)> Callback);
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORM_UTILS_SPLIT_MODULE_BY_CATEGORY_H

llvm/lib/Transforms/Utils/CMakeLists.txt

@@ -84,6 +84,7 @@ add_llvm_component_library(LLVMTransformUtils
   SimplifyLibCalls.cpp
   SizeOpts.cpp
   SplitModule.cpp
+  SplitModuleByCategory.cpp
   StripNonLineTableDebugInfo.cpp
   SymbolRewriter.cpp
   UnifyFunctionExitNodes.cpp

llvm/lib/Transforms/Utils/SplitModuleByCategory.cpp

@@ -0,0 +1,323 @@
+//===-------- SplitModuleByCategory.cpp - split a module by categories ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// See comments in the header.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/SplitModuleByCategory.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+#include <map>
+#include <string>
+#include <utility>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "split-module-by-category"
+
+namespace {
+
+// A vector that contains a group of function with the same category.
+using EntryPointSet = SetVector<const Function *>;
+
+/// Represents a group of functions with one category.
+struct EntryPointGroup {
+  int ID;
+  EntryPointSet Functions;
+
+  EntryPointGroup() = default;
+
+  EntryPointGroup(int ID, EntryPointSet &&Functions = EntryPointSet())
+      : ID(ID), Functions(std::move(Functions)) {}
+
+  void clear() { Functions.clear(); }
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  LLVM_DUMP_METHOD void dump() const {
+    constexpr size_t INDENT = 4;
+    dbgs().indent(INDENT) << "ENTRY POINTS"
+                          << " " << ID << " {\n";
+    for (const Function *F : Functions)
+      dbgs().indent(INDENT) << "  " << F->getName() << "\n";
+
+    dbgs().indent(INDENT) << "}\n";
+  }
+#endif
+};
+
+/// Annotates an llvm::Module with information necessary to perform and track
+/// the result of code (llvm::Module instances) splitting:
+/// - entry points group from the module.
+class ModuleDesc {
+  std::unique_ptr<Module> M;
+  EntryPointGroup EntryPoints;
+
+public:
+  ModuleDesc(std::unique_ptr<Module> M,
+             EntryPointGroup &&EntryPoints = EntryPointGroup())
+      : M(std::move(M)), EntryPoints(std::move(EntryPoints)) {
+    assert(this->M && "Module should be non-null");
+  }
+
+  Module &getModule() { return *M; }
+  const Module &getModule() const { return *M; }
+
+  std::unique_ptr<Module> releaseModule() {
+    EntryPoints.clear();
+    return std::move(M);
+  }
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  LLVM_DUMP_METHOD void dump() const {
+    dbgs() << "ModuleDesc[" << M->getName() << "] {\n";
+    EntryPoints.dump();
+    dbgs() << "}\n";
+  }
+#endif
+};
+
+bool isKernel(const Function &F) {
+  return F.getCallingConv() == CallingConv::SPIR_KERNEL ||
+         F.getCallingConv() == CallingConv::AMDGPU_KERNEL ||
+         F.getCallingConv() == CallingConv::PTX_Kernel;
+}
+
+// Represents "dependency" or "use" graph of global objects (functions and
+// global variables) in a module. It is used during code split to
+// understand which global variables and functions (other than entry points)
+// should be included into a split module.
+//
+// Nodes of the graph represent LLVM's GlobalObjects, edges "A" -> "B" represent
+// the fact that if "A" is included into a module, then "B" should be included
+// as well.
+//
+// Examples of dependencies which are represented in this graph:
+// - Function FA calls function FB
+// - Function FA uses global variable GA
+// - Global variable GA references (initialized with) function FB
+// - Function FA stores address of a function FB somewhere
+//
+// The following cases are treated as dependencies between global objects:
+// 1. Global object A is used by a global object B in any way (store,
+//    bitcast, phi node, call, etc.): "A" -> "B" edge will be added to the
+//    graph;
+// 2. function A performs an indirect call of a function with signature S and
+//    there is a function B with signature S. "A" -> "B" edge will be added to
+//    the graph;
+class DependencyGraph {
+public:
+  using GlobalSet = SmallPtrSet<const GlobalValue *, 16>;
+
+  DependencyGraph(const Module &M) {
+    // Group functions by their signature to handle case (2) described above
+    DenseMap<const FunctionType *, DependencyGraph::GlobalSet>
+        FuncTypeToFuncsMap;
+    for (const Function &F : M.functions()) {
+      // Kernels can't be called (either directly or indirectly).
+      if (isKernel(F))
+        continue;
+
+      FuncTypeToFuncsMap[F.getFunctionType()].insert(&F);
+    }
+
+    for (const Function &F : M.functions()) {
+      // case (1), see comment above the class definition
+      for (const Value *U : F.users())
+        addUserToGraphRecursively(cast<const User>(U), &F);
+
+      // case (2), see comment above the class definition
+      for (const Instruction &I : instructions(F)) {
+        const CallBase *CB = dyn_cast<CallBase>(&I);
+        if (!CB || !CB->isIndirectCall()) // Direct calls were handled above
+          continue;
+
+        const FunctionType *Signature = CB->getFunctionType();
+        GlobalSet &PotentialCallees = FuncTypeToFuncsMap[Signature];
+        Graph[&F].insert(PotentialCallees.begin(), PotentialCallees.end());
+      }
+    }
+
+    // And every global variable (but their handling is a bit simpler)
+    for (const GlobalVariable &GV : M.globals())
+      for (const Value *U : GV.users())
+        addUserToGraphRecursively(cast<const User>(U), &GV);
+  }
+
+  iterator_range<GlobalSet::const_iterator>
+  dependencies(const GlobalValue *Val) const {
+    auto It = Graph.find(Val);
+    return (It == Graph.end())
+               ? make_range(EmptySet.begin(), EmptySet.end())
+               : make_range(It->second.begin(), It->second.end());
+  }
+
+private:
+  void addUserToGraphRecursively(const User *Root, const GlobalValue *V) {
+    SmallVector<const User *, 8> WorkList;
+    WorkList.push_back(Root);
+
+    while (!WorkList.empty()) {
+      const User *U = WorkList.pop_back_val();
+      if (const auto *I = dyn_cast<const Instruction>(U)) {
+        const Function *UFunc = I->getFunction();
+        Graph[UFunc].insert(V);
+      } else if (isa<const Constant>(U)) {
+        if (const auto *GV = dyn_cast<const GlobalVariable>(U))
+          Graph[GV].insert(V);
+        // This could be a global variable or some constant expression (like
+        // bitcast or gep). We trace users of this constant further to reach
+        // global objects they are used by and add them to the graph.
+        for (const User *UU : U->users())
+          WorkList.push_back(UU);
+      } else {
+        llvm_unreachable("Unhandled type of function user");
+      }
+    }
+  }
+
+  DenseMap<const GlobalValue *, GlobalSet> Graph;
+  SmallPtrSet<const GlobalValue *, 1> EmptySet;
+};
+
+void collectFunctionsAndGlobalVariablesToExtract(
+    SetVector<const GlobalValue *> &GVs, const Module &M,
+    const EntryPointGroup &ModuleEntryPoints, const DependencyGraph &DG) {
+  // We start with module entry points
+  for (const Function *F : ModuleEntryPoints.Functions)
+    GVs.insert(F);
+
+  // Non-discardable global variables are also include into the initial set
+  for (const GlobalVariable &GV : M.globals())
+    if (!GV.isDiscardableIfUnused())
+      GVs.insert(&GV);
+
+  // GVs has SetVector type. This type inserts a value only if it is not yet
+  // present there. So, recursion is not expected here.
+  size_t Idx = 0;
+  while (Idx < GVs.size()) {
+    const GlobalValue *Obj = GVs[Idx++];
+
+    for (const GlobalValue *Dep : DG.dependencies(Obj)) {
+      if (const auto *Func = dyn_cast<const Function>(Dep)) {
+        if (!Func->isDeclaration())
+          GVs.insert(Func);
+      } else {
+        GVs.insert(Dep); // Global variables are added unconditionally
+      }
+    }
+  }
+}
+
+ModuleDesc extractSubModule(const Module &M,
+                            const SetVector<const GlobalValue *> &GVs,
+                            EntryPointGroup &&ModuleEntryPoints) {
+  ValueToValueMapTy VMap;
+  // Clone definitions only for needed globals. Others will be added as
+  // declarations and removed later.
+  std::unique_ptr<Module> SubM = CloneModule(
+      M, VMap, [&](const GlobalValue *GV) { return GVs.contains(GV); });
+  // Replace entry points with cloned ones.
+  EntryPointSet NewEPs;
+  const EntryPointSet &EPs = ModuleEntryPoints.Functions;
+  llvm::for_each(
+      EPs, [&](const Function *F) { NewEPs.insert(cast<Function>(VMap[F])); });
+  ModuleEntryPoints.Functions = std::move(NewEPs);
+  return ModuleDesc{std::move(SubM), std::move(ModuleEntryPoints)};
+}
+
+// The function produces a copy of input LLVM IR module M with only those
+// functions and globals that can be called from entry points that are specified
+// in ModuleEntryPoints vector, in addition to the entry point functions.
+ModuleDesc extractCallGraph(const Module &M,
+                            EntryPointGroup &&ModuleEntryPoints,
+                            const DependencyGraph &DG) {
+  SetVector<const GlobalValue *> GVs;
+  collectFunctionsAndGlobalVariablesToExtract(GVs, M, ModuleEntryPoints, DG);
+
+  ModuleDesc SplitM = extractSubModule(M, GVs, std::move(ModuleEntryPoints));
+  LLVM_DEBUG(SplitM.dump());
+  return SplitM;
+}
+
+using EntryPointGroupVec = SmallVector<EntryPointGroup>;
+
+/// Module Splitter.
+/// It gets a module and a collection of entry points groups.
+/// Each group specifies subset entry points from input module that should be
+/// included in a split module.
+class ModuleSplitter {
+private:
+  std::unique_ptr<Module> M;
+  EntryPointGroupVec Groups;
+  DependencyGraph DG;
+
+private:
+  EntryPointGroup drawEntryPointGroup() {
+    assert(Groups.size() > 0 && "Reached end of entry point groups list.");
+    EntryPointGroup Group = std::move(Groups.back());
+    Groups.pop_back();
+    return Group;
+  }
+
+public:
+  ModuleSplitter(std::unique_ptr<Module> Module, EntryPointGroupVec &&GroupVec)
+      : M(std::move(Module)), Groups(std::move(GroupVec)), DG(*M) {
+    assert(!Groups.empty() && "Entry points groups collection is empty!");
+  }
+
+  /// Gets next subsequence of entry points in an input module and provides
+  /// split submodule containing these entry points and their dependencies.
+  ModuleDesc getNextSplit() {
+    return extractCallGraph(*M, drawEntryPointGroup(), DG);
+  }
+
+  /// Check that there are still submodules to split.
+  bool hasMoreSplits() const { return Groups.size() > 0; }
+};
+
+EntryPointGroupVec selectEntryPointGroups(
+    const Module &M, function_ref<std::optional<int>(const Function &F)> EPC) {
+  // std::map is used here to ensure stable ordering of entry point groups,
+  // which is based on their contents, this greatly helps LIT tests
+  // Note: EPC is allowed to return big identifiers. Therefore, we use
+  // std::map + SmallVector approach here.
+  std::map<int, EntryPointSet> EntryPointsMap;
+
+  for (const auto &F : M.functions())
+    if (std::optional<int> Category = EPC(F); Category)
+      EntryPointsMap[*Category].insert(&F);
+
+  EntryPointGroupVec Groups;
+  Groups.reserve(EntryPointsMap.size());
+  for (auto &[Key, EntryPoints] : EntryPointsMap)
+    Groups.emplace_back(Key, std::move(EntryPoints));
+
+  return Groups;
+}
+
+} // namespace
+
+void llvm::splitModuleTransitiveFromEntryPoints(
+    std::unique_ptr<Module> M,
+    function_ref<std::optional<int>(const Function &F)> EntryPointCategorizer,
+    function_ref<void(std::unique_ptr<Module> Part)> Callback) {
+  EntryPointGroupVec Groups = selectEntryPointGroups(*M, EntryPointCategorizer);
+  ModuleSplitter Splitter(std::move(M), std::move(Groups));
+  while (Splitter.hasMoreSplits()) {
+    ModuleDesc MD = Splitter.getNextSplit();
+    Callback(std::move(MD.releaseModule()));
+  }
+}

llvm/test/tools/llvm-split/SplitByCategory/amd-kernel-split.ll

@@ -0,0 +1,17 @@
+; -- Per-kernel split
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %tC
+; RUN: FileCheck %s -input-file=%tC_0.ll --check-prefixes CHECK-A0
+; RUN: FileCheck %s -input-file=%tC_1.ll --check-prefixes CHECK-A1
+
+define dso_local amdgpu_kernel void @KernelA() {
+  ret void
+}
+
+define dso_local amdgpu_kernel void @KernelB() {
+  ret void
+}
+
+; CHECK-A0: define dso_local amdgpu_kernel void @KernelB()
+; CHECK-A0-NOT: define dso_local amdgpu_kernel void @KernelA()
+; CHECK-A1-NOT: define dso_local amdgpu_kernel void @KernelB()
+; CHECK-A1: define dso_local amdgpu_kernel void @KernelA()

llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain1.ll

@@ -0,0 +1,75 @@
+; Check that Module splitting can trace through more complex call stacks
+; involving several nested indirect calls.
+
+; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
+; RUN:     --implicit-check-not @foo --implicit-check-not @kernel_A \
+; RUN:     --implicit-check-not @kernel_B --implicit-check-not @baz
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix CHECK1 \
+; RUN:     --implicit-check-not @kernel_A --implicit-check-not @kernel_C
+; RUN: FileCheck %s -input-file=%t_2.ll --check-prefix CHECK2 \
+; RUN:     --implicit-check-not @foo --implicit-check-not @bar \
+; RUN:     --implicit-check-not @BAZ --implicit-check-not @kernel_B \
+; RUN:     --implicit-check-not @kernel_C
+
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
+; RUN:     --implicit-check-not @foo --implicit-check-not @kernel_A \
+; RUN:     --implicit-check-not @kernel_B
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix CHECK1 \
+; RUN:     --implicit-check-not @kernel_A --implicit-check-not @kernel_C
+; RUN: FileCheck %s -input-file=%t_2.ll --check-prefix CHECK2 \
+; RUN:     --implicit-check-not @foo --implicit-check-not @bar \
+; RUN:     --implicit-check-not @BAZ --implicit-check-not @kernel_B \
+; RUN:     --implicit-check-not @kernel_C
+
+; CHECK0-DAG: define spir_kernel void @kernel_C
+; CHECK0-DAG: define spir_func i32 @bar
+; CHECK0-DAG: define spir_func void @baz
+; CHECK0-DAG: define spir_func void @BAZ
+
+; CHECK1-DAG: define spir_kernel void @kernel_B
+; CHECK1-DAG: define {{.*}}spir_func i32 @foo
+; CHECK1-DAG: define spir_func i32 @bar
+; CHECK1-DAG: define spir_func void @baz
+; CHECK1-DAG: define spir_func void @BAZ
+
+; CHECK2-DAG: define spir_kernel void @kernel_A
+; CHECK2-DAG: define {{.*}}spir_func void @baz
+
+define spir_func i32 @foo(i32 (i32, void ()*)* %ptr1, void ()* %ptr2) {
+  %1 = call spir_func i32 %ptr1(i32 42, void ()* %ptr2)
+  ret i32 %1
+}
+
+define spir_func i32 @bar(i32 %arg, void ()* %ptr) {
+  call spir_func void %ptr()
+  ret i32 %arg
+}
+
+define spir_func void @baz() {
+  ret void
+}
+
+define spir_func void @BAZ() {
+  ret void
+}
+
+define spir_kernel void @kernel_A() #0 {
+  call spir_func void @baz()
+  ret void
+}
+
+define spir_kernel void @kernel_B() #1 {
+  call spir_func i32 @foo(i32 (i32, void ()*)* null, void ()* null)
+  ret void
+}
+
+define spir_kernel void @kernel_C() #2 {
+  call spir_func i32 @bar(i32 42, void ()* null)
+  ret void
+}
+
+attributes #0 = { "module-id"="TU1.cpp" }
+attributes #1 = { "module-id"="TU2.cpp" }
+attributes #2 = { "module-id"="TU3.cpp" }

llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain2.ll

@@ -0,0 +1,49 @@
+; Check that Module splitting can trace indirect calls through signatures.
+
+; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
+; RUN:     --implicit-check-not @kernel_A --implicit-check-not @bbb
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix CHECK1 \
+; RUN:     --implicit-check-not @kernel_B --implicit-check-not @ccc
+
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
+; RUN:     --implicit-check-not @kernel_A
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix CHECK1 \
+; RUN:     --implicit-check-not @kernel_B
+
+; CHECK0-DAG: define spir_kernel void @kernel_B
+; CHECK0-DAG: define spir_func void @aaa(i32 %0, i32 %1)
+; CHECK0-DAG: define spir_func void @ccc(ptr %ptr)
+
+; CHECK1-DAG: define spir_kernel void @kernel_A
+; CHECK1-DAG: define spir_func void @aaa(i32 %0, i32 %1)
+; CHECK1-DAG: define spir_func void @bbb(ptr %ptr)
+
+
+define spir_func void @aaa(i32 %0, i32 %1) {
+  ret void
+}
+
+define spir_func void @bbb(void (i32, i32)* %ptr) {
+  call spir_func void %ptr(i32 0, i32 0)
+  ret void
+}
+
+define spir_func void @ccc(void (i32, i32)* %ptr) {
+  call spir_func void %ptr(i32 0, i32 0)
+  ret void
+}
+
+define spir_kernel void @kernel_A() #0 {
+  call spir_func void @bbb(void (i32, i32)* null)
+  ret void
+}
+
+define spir_kernel void @kernel_B() #1 {
+  call spir_func void @ccc(void (i32, i32)* null)
+  ret void
+}
+
+attributes #0 = { "module-id"="TU1.cpp" }
+attributes #1 = { "module-id"="TU2.cpp" }

llvm/test/tools/llvm-split/SplitByCategory/module-split-func-ptr.ll

@@ -0,0 +1,38 @@
+; This test checks that Module splitting can properly perform device code split by tracking
+; all uses of functions (not only direct calls).
+
+; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix=CHECK-IR0
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix=CHECK-IR1
+
+; CHECK-IR0: define dso_local spir_kernel void @kernelA
+;
+; CHECK-IR1: @FuncTable = weak global ptr @func
+; CHECK-IR1: define {{.*}} i32 @func
+; CHECK-IR1: define weak_odr dso_local spir_kernel void @kernelB
+
+@FuncTable = weak global ptr @func, align 8
+
+define dso_local spir_func i32 @func(i32 %a) {
+entry:
+  ret i32 %a
+}
+
+define weak_odr dso_local spir_kernel void @kernelB() #0 {
+entry:
+  %0 = call i32 @indirect_call(ptr addrspace(4) addrspacecast ( ptr getelementptr inbounds ( [1 x ptr] , ptr @FuncTable, i64 0, i64 0) to ptr addrspace(4)), i32 0)
+  ret void
+}
+
+define dso_local spir_kernel void @kernelA() #1 {
+entry:
+  ret void
+}
+
+declare dso_local spir_func i32 @indirect_call(ptr addrspace(4), i32) local_unnamed_addr
+
+attributes #0 = { "module-id"="TU1.cpp" }
+attributes #1 = { "module-id"="TU2.cpp" }
+
+; CHECK: kernel1
+; CHECK: kernel2

llvm/test/tools/llvm-split/SplitByCategory/one-kernel-per-module.ll

@@ -0,0 +1,93 @@
+; Test checks "kernel" splitting mode.
+
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %t.files
+; RUN: FileCheck %s -input-file=%t.files_0.ll --check-prefixes CHECK-MODULE0,CHECK
+; RUN: FileCheck %s -input-file=%t.files_1.ll --check-prefixes CHECK-MODULE1,CHECK
+; RUN: FileCheck %s -input-file=%t.files_2.ll --check-prefixes CHECK-MODULE2,CHECK
+
+;CHECK-MODULE0: @GV = internal addrspace(1) constant [1 x i32] [i32 42], align 4
+;CHECK-MODULE1-NOT: @GV
+;CHECK-MODULE2-NOT: @GV
+@GV = internal addrspace(1) constant [1 x i32] [i32 42], align 4
+
+; CHECK-MODULE0-NOT: define dso_local spir_kernel void @TU0_kernelA
+; CHECK-MODULE1-NOT: define dso_local spir_kernel void @TU0_kernelA
+; CHECK-MODULE2: define dso_local spir_kernel void @TU0_kernelA
+define dso_local spir_kernel void @TU0_kernelA() #0 {
+entry:
+; CHECK-MODULE2: call spir_func void @foo()
+  call spir_func void @foo()
+  ret void
+}
+
+; CHECK-MODULE0-NOT: define {{.*}} spir_func void @foo()
+; CHECK-MODULE1-NOT: define {{.*}} spir_func void @foo()
+; CHECK-MODULE2: define {{.*}} spir_func void @foo()
+define dso_local spir_func void @foo() {
+entry:
+; CHECK-MODULE2: call spir_func void @bar()
+  call spir_func void @bar()
+  ret void
+}
+
+; CHECK-MODULE0-NOT: define {{.*}} spir_func void @bar()
+; CHECK-MODULE1-NOT: define {{.*}} spir_func void @bar()
+; CHECK-MODULE2: define {{.*}} spir_func void @bar()
+define linkonce_odr dso_local spir_func void @bar() {
+entry:
+  ret void
+}
+
+; CHECK-MODULE0-NOT: define dso_local spir_kernel void @TU0_kernelB()
+; CHECK-MODULE1: define dso_local spir_kernel void @TU0_kernelB()
+; CHECK-MODULE2-NOT: define dso_local spir_kernel void @TU0_kernelB()
+define dso_local spir_kernel void @TU0_kernelB() #0 {
+entry:
+; CHECK-MODULE1: call spir_func void @foo1()
+  call spir_func void @foo1()
+  ret void
+}
+
+; CHECK-MODULE0-NOT: define {{.*}} spir_func void @foo1()
+; CHECK-MODULE1: define {{.*}} spir_func void @foo1()
+; CHECK-MODULE2-NOT: define {{.*}} spir_func void @foo1()
+define dso_local spir_func void @foo1() {
+entry:
+  ret void
+}
+
+; CHECK-MODULE0: define dso_local spir_kernel void @TU1_kernel()
+; CHECK-MODULE1-NOT: define dso_local spir_kernel void @TU1_kernel()
+; CHECK-MODULE2-NOT: define dso_local spir_kernel void @TU1_kernel()
+define dso_local spir_kernel void @TU1_kernel() #1 {
+entry:
+; CHECK-MODULE0: call spir_func void @foo2()
+  call spir_func void @foo2()
+  ret void
+}
+
+; CHECK-MODULE0: define {{.*}} spir_func void @foo2()
+; CHECK-MODULE1-NOT: define {{.*}} spir_func void @foo2()
+; CHECK-MODULE2-NOT: define {{.*}} spir_func void @foo2()
+define dso_local spir_func void @foo2() {
+entry:
+; CHECK-MODULE0: %0 = load i32, ptr addrspace(4) addrspacecast (ptr addrspace(1) @GV to ptr addrspace(4)), align 4
+  %0 = load i32, ptr addrspace(4) getelementptr inbounds ([1 x i32], ptr addrspace(4) addrspacecast (ptr addrspace(1) @GV to ptr addrspace(4)), i64 0, i64 0), align 4
+  ret void
+}
+
+attributes #0 = { "module-id"="TU1.cpp" }
+attributes #1 = { "module-id"="TU2.cpp" }
+
+; Metadata is saved in both modules.
+; CHECK: !opencl.spir.version = !{!0, !0}
+; CHECK: !spirv.Source = !{!1, !1}
+
+!opencl.spir.version = !{!0, !0}
+!spirv.Source = !{!1, !1}
+
+; CHECK; !0 = !{i32 1, i32 2}
+; CHECK; !1 = !{i32 4, i32 100000}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}

llvm/test/tools/llvm-split/SplitByCategory/ptx-kernel-split.ll

@@ -0,0 +1,17 @@
+; -- Per-kernel split
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %tC
+; RUN: FileCheck %s -input-file=%tC_0.ll --check-prefixes CHECK-A0
+; RUN: FileCheck %s -input-file=%tC_1.ll --check-prefixes CHECK-A1
+
+define dso_local ptx_kernel void @KernelA() {
+  ret void
+}
+
+define dso_local ptx_kernel void @KernelB() {
+  ret void
+}
+
+; CHECK-A0: define dso_local ptx_kernel void @KernelB()
+; CHECK-A0-NOT: define dso_local ptx_kernel void @KernelA()
+; CHECK-A1-NOT: define dso_local ptx_kernel void @KernelB()
+; CHECK-A1: define dso_local ptx_kernel void @KernelA()

llvm/test/tools/llvm-split/SplitByCategory/split-by-source.ll

@@ -0,0 +1,88 @@
+; Test checks that kernels are being split by attached module-id metadata and
+; used functions are being moved with kernels that use them.
+
+; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefixes CHECK-TU0,CHECK
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefixes CHECK-TU1,CHECK
+
+; CHECK-TU1-NOT: @GV
+; CHECK-TU0: @GV = internal addrspace(1) constant [1 x i32] [i32 42], align 4
+@GV = internal addrspace(1) constant [1 x i32] [i32 42], align 4
+
+; CHECK-TU0-NOT: define dso_local spir_kernel void @TU1_kernelA
+; CHECK-TU1: define dso_local spir_kernel void @TU1_kernelA
+define dso_local spir_kernel void @TU1_kernelA() #0 {
+entry:
+; CHECK-TU1: call spir_func void @func1_TU1()
+  call spir_func void @func1_TU1()
+  ret void
+}
+
+; CHECK-TU0-NOT: define {{.*}} spir_func void @func1_TU1()
+; CHECK-TU1: define {{.*}} spir_func void @func1_TU1()
+define dso_local spir_func void @func1_TU1() {
+entry:
+; CHECK-TU1: call spir_func void @func2_TU1()
+  call spir_func void @func2_TU1()
+  ret void
+}
+
+; CHECK-TU0-NOT: define {{.*}} spir_func void @func2_TU1()
+; CHECK-TU1: define {{.*}} spir_func void @func2_TU1()
+define linkonce_odr dso_local spir_func void @func2_TU1() {
+entry:
+  ret void
+}
+
+; CHECK-TU0-NOT: define dso_local spir_kernel void @TU1_kernelB()
+; CHECK-TU1: define dso_local spir_kernel void @TU1_kernelB()
+define dso_local spir_kernel void @TU1_kernelB() #0 {
+entry:
+; CHECK-TU1: call spir_func void @func3_TU1()
+  call spir_func void @func3_TU1()
+  ret void
+}
+
+; CHECK-TU0-NOT: define {{.*}} spir_func void @func3_TU1()
+; CHECK-TU1: define {{.*}} spir_func void @func3_TU1()
+define dso_local spir_func void @func3_TU1() {
+entry:
+  ret void
+}
+
+; CHECK-TU0-TXT: TU0_kernel
+; CHECK-TU1-TXT-NOT: TU0_kernel
+
+; CHECK-TU0: define dso_local spir_kernel void @TU0_kernel()
+; CHECK-TU1-NOT: define dso_local spir_kernel void @TU0_kernel()
+define dso_local spir_kernel void @TU0_kernel() #1 {
+entry:
+; CHECK-TU0: call spir_func void @func_TU0()
+  call spir_func void @func_TU0()
+  ret void
+}
+
+; CHECK-TU0: define {{.*}} spir_func void @func_TU0()
+; CHECK-TU1-NOT: define {{.*}} spir_func void @func_TU0()
+define dso_local spir_func void @func_TU0() {
+entry:
+; CHECK-TU0: %0 = load i32, ptr addrspace(4) addrspacecast (ptr addrspace(1) @GV to ptr addrspace(4)), align 4
+  %0 = load i32, ptr addrspace(4) getelementptr inbounds ([1 x i32], ptr addrspace(4) addrspacecast (ptr addrspace(1) @GV to ptr addrspace(4)), i64 0, i64 0), align 4
+  ret void
+}
+
+attributes #0 = { "module-id"="TU1.cpp" }
+attributes #1 = { "module-id"="TU2.cpp" }
+
+; Metadata is saved in both modules.
+; CHECK: !opencl.spir.version = !{!0, !0}
+; CHECK: !spirv.Source = !{!1, !1}
+
+!opencl.spir.version = !{!0, !0}
+!spirv.Source = !{!1, !1}
+
+; CHECK: !0 = !{i32 1, i32 2}
+; CHECK: !1 = !{i32 4, i32 100000}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}

llvm/test/tools/llvm-split/SplitByCategory/split-with-kernel-declarations.ll

@@ -0,0 +1,52 @@
+; The test checks that Module splitting does not treat declarations as entry points.
+
+; RUN: llvm-split -split-by-category=module-id -S < %s -o %t1
+; RUN: FileCheck %s -input-file=%t1_0.ll --check-prefix CHECK-MODULE-ID0
+; RUN: FileCheck %s -input-file=%t1_1.ll --check-prefix CHECK-MODULE-ID1
+
+; RUN: llvm-split -split-by-category=kernel -S < %s -o %t2
+; RUN: FileCheck %s -input-file=%t2_0.ll --check-prefix CHECK-PER-KERNEL0
+; RUN: FileCheck %s -input-file=%t2_1.ll --check-prefix CHECK-PER-KERNEL1
+; RUN: FileCheck %s -input-file=%t2_2.ll --check-prefix CHECK-PER-KERNEL2
+
+; With module-id split, there should be two modules
+; CHECK-MODULE-ID0-NOT: TU0
+; CHECK-MODULE-ID0-NOT: TU1_kernel1
+; CHECK-MODULE-ID0: TU1_kernel0
+;
+; CHECK-MODULE-ID1-NOT: TU1
+; CHECK-MODULE-ID1: TU0_kernel0
+; CHECK-MODULE-ID1: TU0_kernel1
+
+; With per-kernel split, there should be three modules.
+; CHECK-PER-KERNEL0-NOT: TU0
+; CHECK-PER-KERNEL0-NOT: TU1_kernel1
+; CHECK-PER-KERNEL0: TU1_kernel0
+;
+; CHECK-PER-KERNEL1-NOT: TU0_kernel0
+; CHECK-PER-KERNEL1-NOT: TU1
+; CHECK-PER-KERNEL1: TU0_kernel1
+;
+; CHECK-PER-KERNEL2-NOT: TU0_kernel1
+; CHECK-PER-KERNEL2-NOT: TU1
+; CHECK-PER-KERNEL2: TU0_kernel0
+
+
+define spir_kernel void @TU0_kernel0() #0 {
+entry:
+  ret void
+}
+
+define spir_kernel void @TU0_kernel1() #0 {
+entry:
+  ret void
+}
+
+define spir_kernel void @TU1_kernel0() #1 {
+  ret void
+}
+
+declare spir_kernel void @TU1_kernel1() #1
+
+attributes #0 = { "module-id"="TU1.cpp" }
+attributes #1 = { "module-id"="TU2.cpp" }

llvm/tools/llvm-split/CMakeLists.txt

@@ -12,6 +12,7 @@ set(LLVM_LINK_COMPONENTS
   Support
   Target
   TargetParser
+  ipo
   )
 
 add_llvm_tool(llvm-split

llvm/tools/llvm-split/llvm-split.cpp

@@ -11,14 +11,19 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/PassInstrumentation.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
@@ -27,7 +32,9 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/Triple.h"
+#include "llvm/Transforms/IPO/GlobalDCE.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
+#include "llvm/Transforms/Utils/SplitModuleByCategory.h"
 
 using namespace llvm;
 
@@ -70,6 +77,164 @@ static cl::opt<std::string>
     MCPU("mcpu", cl::desc("Target CPU, ignored if --mtriple is not used"),
          cl::value_desc("cpu"), cl::cat(SplitCategory));
 
+enum class SplitByCategoryType {
+  SBCT_ByModuleId,
+  SBCT_ByKernel,
+  SBCT_None,
+};
+
+static cl::opt<SplitByCategoryType> SplitByCategory(
+    "split-by-category",
+    cl::desc("Split by category. If present, splitting by category is used "
+             "with the specified categorization type."),
+    cl::Optional, cl::init(SplitByCategoryType::SBCT_None),
+    cl::values(clEnumValN(SplitByCategoryType::SBCT_ByModuleId, "module-id",
+                          "one output module per translation unit marked with "
+                          "\"module-id\" attribute"),
+               clEnumValN(SplitByCategoryType::SBCT_ByKernel, "kernel",
+                          "one output module per kernel")),
+    cl::cat(SplitCategory));
+
+static cl::opt<bool> OutputAssembly{
+    "S", cl::desc("Write output as LLVM assembly"), cl::cat(SplitCategory)};
+
+void writeStringToFile(StringRef Content, StringRef Path) {
+  std::error_code EC;
+  raw_fd_ostream OS(Path, EC);
+  if (EC) {
+    errs() << formatv("error opening file: {0}, error: {1}\n", Path,
+                      EC.message());
+    exit(1);
+  }
+
+  OS << Content << "\n";
+}
+
+void writeModuleToFile(const Module &M, StringRef Path, bool OutputAssembly) {
+  int FD = -1;
+  if (std::error_code EC = sys::fs::openFileForWrite(Path, FD)) {
+    errs() << formatv("error opening file: {0}, error: {1}", Path, EC.message())
+           << '\n';
+    exit(1);
+  }
+
+  raw_fd_ostream OS(FD, /*ShouldClose*/ true);
+  if (OutputAssembly)
+    M.print(OS, /*AssemblyAnnotationWriter*/ nullptr);
+  else
+    WriteBitcodeToFile(M, OS);
+}
+
+/// EntryPointCategorizer is used for splitting by category either by module-id
+/// or by kernels. It doesn't provide categories for functions other than
+/// kernels. Categorizer computes a string key for the given Function and
+/// records the association between the string key and an integer category. If a
+/// string key is already belongs to some category than the corresponding
+/// integer category is returned.
+class EntryPointCategorizer {
+public:
+  EntryPointCategorizer(SplitByCategoryType Type) : Type(Type) {}
+
+  EntryPointCategorizer() = delete;
+  EntryPointCategorizer(EntryPointCategorizer &) = delete;
+  EntryPointCategorizer &operator=(const EntryPointCategorizer &) = delete;
+  EntryPointCategorizer(EntryPointCategorizer &&) = default;
+  EntryPointCategorizer &operator=(EntryPointCategorizer &&) = default;
+
+  /// Returns integer specifying the category for the given \p F.
+  /// If the given function isn't a kernel then returns std::nullopt.
+  std::optional<int> operator()(const Function &F) {
+    if (!isEntryPoint(F))
+      return std::nullopt; // skip the function.
+
+    auto StringKey = computeFunctionCategory(Type, F);
+    if (auto it = StrKeyToID.find(StringRef(StringKey)); it != StrKeyToID.end())
+      return it->second;
+
+    int ID = static_cast<int>(StrKeyToID.size());
+    return StrKeyToID.try_emplace(std::move(StringKey), ID).first->second;
+  }
+
+private:
+  static bool isEntryPoint(const Function &F) {
+    if (F.isDeclaration())
+      return false;
+
+    return F.getCallingConv() == CallingConv::SPIR_KERNEL ||
+           F.getCallingConv() == CallingConv::AMDGPU_KERNEL ||
+           F.getCallingConv() == CallingConv::PTX_Kernel;
+  }
+
+  static SmallString<0> computeFunctionCategory(SplitByCategoryType Type,
+                                                const Function &F) {
+    static constexpr char ATTR_MODULE_ID[] = "module-id";
+    SmallString<0> Key;
+    switch (Type) {
+    case SplitByCategoryType::SBCT_ByKernel:
+      Key = F.getName().str();
+      break;
+    case SplitByCategoryType::SBCT_ByModuleId:
+      Key = F.getFnAttribute(ATTR_MODULE_ID).getValueAsString().str();
+      break;
+    default:
+      llvm_unreachable("unexpected mode.");
+    }
+
+    return Key;
+  }
+
+private:
+  struct KeyInfo {
+    static SmallString<0> getEmptyKey() { return SmallString<0>(""); }
+
+    static SmallString<0> getTombstoneKey() { return SmallString<0>("-"); }
+
+    static bool isEqual(const SmallString<0> &LHS, const SmallString<0> &RHS) {
+      return LHS == RHS;
+    }
+
+    static unsigned getHashValue(const SmallString<0> &S) {
+      return llvm::hash_value(StringRef(S));
+    }
+  };
+
+  SplitByCategoryType Type;
+  DenseMap<SmallString<0>, int, KeyInfo> StrKeyToID;
+};
+
+void cleanupModule(Module &M) {
+  ModuleAnalysisManager MAM;
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+  ModulePassManager MPM;
+  MPM.addPass(GlobalDCEPass()); // Delete unreachable globals.
+  MPM.run(M, MAM);
+}
+
+Error runSplitModuleByCategory(std::unique_ptr<Module> M) {
+  size_t OutputID = 0;
+  auto PostSplitCallback = [&](std::unique_ptr<Module> MPart) {
+    if (verifyModule(*MPart)) {
+      errs() << "Broken Module!\n";
+      exit(1);
+    }
+
+    // TODO: DCE is a crucial pass since it removes unused declarations.
+    //       At the moment, LIT checking can't be perfomed without DCE.
+    cleanupModule(*MPart);
+    size_t ID = OutputID;
+    ++OutputID;
+    StringRef ModuleSuffix = OutputAssembly ? ".ll" : ".bc";
+    std::string ModulePath =
+        (Twine(OutputFilename) + "_" + Twine(ID) + ModuleSuffix).str();
+    writeModuleToFile(*MPart, ModulePath, OutputAssembly);
+  };
+
+  auto Categorizer = EntryPointCategorizer(SplitByCategory);
+  splitModuleTransitiveFromEntryPoints(std::move(M), Categorizer,
+                                       PostSplitCallback);
+  return Error::success();
+}
+
 int main(int argc, char **argv) {
   InitLLVM X(argc, argv);
 
@@ -123,6 +288,17 @@ int main(int argc, char **argv) {
     Out->keep();
   };
 
+  if (SplitByCategory != SplitByCategoryType::SBCT_None) {
+    auto E = runSplitModuleByCategory(std::move(M));
+    if (E) {
+      errs() << E << "\n";
+      Err.print(argv[0], errs());
+      return 1;
+    }
+
+    return 0;
+  }
+
   if (TM) {
     if (PreserveLocals) {
       errs() << "warning: --preserve-locals has no effect when using "