[offload][SYCL] Add SYCL Module splitting

llvm/include/llvm/Transforms/Utils/SYCLSplitModule.h

@@ -10,8 +10,8 @@
 // of the split is new modules containing corresponding callgraph.
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_SYCL_SPLIT_MODULE_H
-#define LLVM_SYCL_SPLIT_MODULE_H
+#ifndef LLVM_TRANSFORMS_UTILS_SYCLSPLITMODULE_H
+#define LLVM_TRANSFORMS_UTILS_SYCLSPLITMODULE_H
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -68,4 +68,4 @@ SYCLSplitModule(std::unique_ptr<Module> M, ModuleSplitterSettings Settings);
 
 } // namespace llvm
 
-#endif // LLVM_SYCL_SPLIT_MODULE_H
+#endif // LLVM_TRANSFORMS_UTILS_SYCLSPLITMODULE_H

llvm/include/llvm/Transforms/Utils/SYCLUtils.h

@@ -10,14 +10,14 @@
 #ifndef LLVM_TRANSFORMS_UTILS_SYCLUTILS_H
 #define LLVM_TRANSFORMS_UTILS_SYCLUTILS_H
 
-#include <string>
-#include <vector>
+#include <llvm/ADT/SmallString.h>
+#include <llvm/ADT/SmallVector.h>
 
 namespace llvm {
 
 class raw_ostream;
 
-using SYCLStringTable = std::vector<std::vector<std::string>>;
+using SYCLStringTable = SmallVector<SmallVector<SmallString<64>>>;
 
 void writeSYCLStringTable(const SYCLStringTable &Table, raw_ostream &OS);
 

llvm/lib/Transforms/Utils/SYCLSplitModule.cpp

@@ -1,4 +1,4 @@
-//===-------- SYCLSplitModule.cpp - split a module into callgraphs --------===//
+//===-------- SYCLSplitModule.cpp - Split a module into call graphs -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

llvm/lib/Transforms/Utils/SYCLUtils.cpp

@@ -15,11 +15,10 @@ namespace llvm {
 
 void writeSYCLStringTable(const SYCLStringTable &Table, raw_ostream &OS) {
   assert(Table.size() > 0 && "table should contain at least column titles");
-  size_t numberColumns = Table[0].size();
-  assert(numberColumns > 0 && "table should be non-empty");
+  assert(Table[0].size() > 0 && "table should be non-empty");
   OS << '[' << join(Table[0].begin(), Table[0].end(), "|") << "]\n";
   for (size_t I = 1, E = Table.size(); I != E; ++I) {
-    assert(Table[I].size() == numberColumns && "row's size should be equal");
+    assert(Table[I].size() == Table[0].size() && "row's size should be equal");
     OS << join(Table[I].begin(), Table[I].end(), "|") << '\n';
   }
 }

llvm/test/tools/llvm-split/SYCL/device-code-split/split-by-source.ll

@@ -0,0 +1,107 @@
+; Test checks that kernels are being split by attached TU metadata and
+; used functions are being moved with kernels that use them.
+
+; RUN: llvm-split -sycl-split=source -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
+; RUN: FileCheck %s -input-file=%t_0.sym --check-prefixes CHECK-TU0-TXT
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefixes CHECK-TU1-TXT
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir64-unknown-linux"
+
+; 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-TU1: define dso_local spir_kernel void @TU1_kernel0
+; CHECK-TU1-TXT: TU1_kernel0
+; CHECK-TU0-NOT: define dso_local spir_kernel void @TU1_kernel0
+; CHECK-TU0-TXT-NOT: TU1_kernel0
+
+; CHECK-TU1: call spir_func void @func1_TU1()
+
+define dso_local spir_kernel void @TU1_kernel0() #0 {
+entry:
+  call spir_func void @func1_TU1()
+  ret void
+}
+
+; CHECK-TU1: define {{.*}} spir_func void @func1_TU1()
+; CHECK-TU0-NOT: define {{.*}} spir_func void @func1_TU1()
+
+; CHECK-TU1: call spir_func void @func2_TU1()
+
+define dso_local spir_func void @func1_TU1() {
+entry:
+  call spir_func void @func2_TU1()
+  ret void
+}
+
+; CHECK-TU1: define {{.*}} spir_func void @func2_TU1()
+; CHECK-TU0-NOT: define {{.*}} spir_func void @func2_TU1()
+
+define linkonce_odr dso_local spir_func void @func2_TU1() {
+entry:
+  ret void
+}
+
+; CHECK-TU1: define dso_local spir_kernel void @TU1_kernel1()
+; CHECK-TU1-TXT: TU1_kernel1
+; CHECK-TU0-NOT: define dso_local spir_kernel void @TU1_kernel1()
+; CHECK-TU0-TXT-NOT: TU1_kernel1
+
+; CHECK-TU1: call spir_func void @func3_TU1()
+
+define dso_local spir_kernel void @TU1_kernel1() #0 {
+entry:
+  call spir_func void @func3_TU1()
+  ret void
+}
+
+; CHECK-TU1: define {{.*}} spir_func void @func3_TU1()
+; CHECK-TU0-NOT: define {{.*}} spir_func void @func3_TU1()
+
+define dso_local spir_func void @func3_TU1() {
+entry:
+  ret void
+}
+
+; CHECK-TU1-NOT: define dso_local spir_kernel void @TU0_kernel()
+; CHECK-TU1-TXT-NOT: TU0_kernel
+; CHECK-TU0: define dso_local spir_kernel void @TU0_kernel()
+; CHECK-TU0-TXT: TU0_kernel
+
+; CHECK-TU0: call spir_func void @func_TU0()
+
+define dso_local spir_kernel void @TU0_kernel() #1 {
+entry:
+  call spir_func void @func_TU0()
+  ret void
+}
+
+; CHECK-TU1-NOT: define {{.*}} spir_func void @func_TU0()
+; CHECK-TU0: 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 = { "sycl-module-id"="TU1.cpp" }
+attributes #1 = { "sycl-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/tools/llvm-split/llvm-split.cpp

@@ -11,6 +11,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#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"
@@ -32,9 +34,6 @@
 #include "llvm/Transforms/Utils/SYCLUtils.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 
-#include <string>
-#include <vector>
-
 using namespace llvm;
 
 static cl::OptionCategory SplitCategory("Split Options");
@@ -90,7 +89,7 @@ static cl::opt<IRSplitMode> SYCLSplitMode(
 static cl::opt<bool> OutputAssembly{
     "S", cl::desc("Write output as LLVM assembly"), cl::cat(SplitCategory)};
 
-void writeStringToFile(std::string_view Content, StringRef Path) {
+void writeStringToFile(StringRef Content, StringRef Path) {
   std::error_code EC;
   raw_fd_ostream OS(Path, EC);
   if (EC) {
@@ -103,13 +102,19 @@ void writeStringToFile(std::string_view Content, StringRef Path) {
 
 void writeSplitModulesAsTable(ArrayRef<ModuleAndSYCLMetadata> Modules,
                               StringRef Path) {
-  std::vector<std::string> Columns = {"Code", "Symbols"};
+  SmallVector<SmallString<64>> Columns;
+  Columns.emplace_back("Code");
+  Columns.emplace_back("Symbols");
+
   SYCLStringTable Table;
   Table.emplace_back(std::move(Columns));
   for (const auto &[I, SM] : enumerate(Modules)) {
-    std::string SymbolsFile = (Twine(Path) + "_" + Twine(I) + ".sym").str();
+    SmallString<128> SymbolsFile;
+    (Twine(Path) + "_" + Twine(I) + ".sym").toVector(SymbolsFile);
     writeStringToFile(SM.Symbols, SymbolsFile);
-    std::vector<std::string> Row = {SM.ModuleFilePath, SymbolsFile};
+    SmallVector<SmallString<64>> Row;
+    Row.emplace_back(SM.ModuleFilePath);
+    Row.emplace_back(SymbolsFile);
     Table.emplace_back(std::move(Row));
   }
 
@@ -194,6 +199,7 @@ int main(int argc, char **argv) {
     if (E) {
       errs() << E << "\n";
       Err.print(argv[0], errs());
+      return 1;
     }
 
     return 0;