[SYCL] Implement device library linking for runtime compilation

sycl-jit/jit-compiler/CMakeLists.txt

@@ -23,6 +23,7 @@ add_llvm_library(sycl-jit
    IPO
    TransformUtils
    Passes
+   IRReader
    Linker
    ScalarOpts
    InstCombine

sycl-jit/jit-compiler/lib/KernelFusion.cpp

@@ -244,17 +244,18 @@ extern "C" JITResult compileSYCL(InMemoryFile SourceFile,
     return errorToFusionResult(ModuleOrErr.takeError(),
                                "Device compilation failed");
   }
-  std::unique_ptr<llvm::Module> Module = std::move(*ModuleOrErr);
 
-  SYCLKernelInfo Kernel;
-  auto Error = translation::KernelTranslator::translateKernel(
-      Kernel, *Module, JITContext::getInstance(), BinaryFormat::SPIRV);
+  std::unique_ptr<llvm::LLVMContext> Context;
+  std::unique_ptr<llvm::Module> Module = std::move(*ModuleOrErr);
+  Context.reset(&Module->getContext());
 
-  auto *LLVMCtx = &Module->getContext();
-  Module.reset();
-  delete LLVMCtx;
+  if (auto Error = linkDefaultDeviceLibraries(*Module, UserArgs)) {
+    return errorToFusionResult(std::move(Error), "Device linking failed");
+  }
 
-  if (Error) {
+  SYCLKernelInfo Kernel;
+  if (auto Error = translation::KernelTranslator::translateKernel(
+          Kernel, *Module, JITContext::getInstance(), BinaryFormat::SPIRV)) {
     return errorToFusionResult(std::move(Error), "SPIR-V translation failed");
   }
 

sycl-jit/jit-compiler/lib/rtc/DeviceCompilation.cpp

@@ -15,6 +15,16 @@
 #include <clang/Tooling/CompilationDatabase.h>
 #include <clang/Tooling/Tooling.h>
 
+#include <llvm/IRReader/IRReader.h>
+#include <llvm/Linker/Linker.h>
+
+#include <array>
+
+using namespace clang;
+using namespace clang::tooling;
+using namespace clang::driver;
+using namespace llvm;
+
 #ifdef _GNU_SOURCE
 #include <dlfcn.h>
 static char X; // Dummy symbol, used as an anchor for `dlinfo` below.
@@ -49,9 +59,6 @@ static const std::string &getDPCPPRoot() {
 }
 
 namespace {
-using namespace clang;
-using namespace clang::tooling;
-using namespace clang::driver;
 
 struct GetLLVMModuleAction : public ToolAction {
   // Code adapted from `FrontendActionFactory::runInvocation`.
@@ -96,23 +103,20 @@ struct GetLLVMModuleAction : public ToolAction {
 
 } // anonymous namespace
 
-llvm::Expected<std::unique_ptr<llvm::Module>>
+Expected<std::unique_ptr<llvm::Module>>
 jit_compiler::compileDeviceCode(InMemoryFile SourceFile,
                                 View<InMemoryFile> IncludeFiles,
                                 View<const char *> UserArgs) {
   const std::string &DPCPPRoot = getDPCPPRoot();
   if (DPCPPRoot == InvalidDPCPPRoot) {
-    return llvm::createStringError("Could not locate DPCPP root directory");
+    return createStringError("Could not locate DPCPP root directory");
   }
 
   SmallVector<std::string> CommandLine = {"-fsycl-device-only"};
-  // TODO: Allow instrumentation again when device library linking is
-  //       implemented.
-  CommandLine.push_back("-fno-sycl-instrument-device-code");
   CommandLine.append(UserArgs.begin(), UserArgs.end());
-  clang::tooling::FixedCompilationDatabase DB{".", CommandLine};
+  FixedCompilationDatabase DB{".", CommandLine};
 
-  clang::tooling::ClangTool Tool{DB, {SourceFile.Path}};
+  ClangTool Tool{DB, {SourceFile.Path}};
 
   // Set up in-memory filesystem.
   Tool.mapVirtualFile(SourceFile.Path, SourceFile.Contents);
@@ -143,5 +147,90 @@ jit_compiler::compileDeviceCode(InMemoryFile SourceFile,
   }
 
   // TODO: Capture compiler errors from the ClangTool.
-  return llvm::createStringError("Unable to obtain LLVM module");
+  return createStringError("Unable to obtain LLVM module");
+}
+
+Error jit_compiler::linkDefaultDeviceLibraries(llvm::Module &Module,
+                                               View<const char *> UserArgs) {
+  // This function mimics the device library selection process
+  // `clang::driver::tools::SYCL::getDeviceLibraries`, assuming a SPIR-V target
+  // (no AoT, no third-party GPUs, no native CPU).
+
+  bool DeviceInstrumentationEnabled = true;
+  for (StringRef UA : UserArgs) {
+    // Check instrumentation-related flags (last occurence determines outcome).
+    if (UA == "-fno-sycl-instrument-device-code") {
+      DeviceInstrumentationEnabled = false;
+      continue;
+    }
+    if (UA == "-fsycl-instrument-device-code") {
+      DeviceInstrumentationEnabled = true;
+      continue;
+    }
+
+    // Issue warning for `-fsycl-device-lib` or `-fno-sycl-device-lib`.
+    // TODO: Is it worth supporting these flags? We're using `LinkOnlyNeeded`
+    //       mode anyways!
+    // TODO: If we keep the warning, it must go into the build log, not onto the
+    //       console.
+    // TODO: The FrontendAction emits a warning that these flags are unused. We
+    //       should probably silence that by removing the argument occurence for
+    //       the compilation step.
+    if (UA.contains("sycl-device-lib")) {
+      errs() << "warning: device library selection with '" << UA
+             << "' is ignored\n";
+    }
+
+    // TODO: Presence of `-fsanitize=address` would require linking
+    //       `libsycl-sanitizer`, but currently compilation fails earlier.
+    assert(!UA.contains("-fsanitize=address") && "Device ASAN unsupported");
+  }
+
+  const std::string &DPCPPRoot = getDPCPPRoot();
+  if (DPCPPRoot == InvalidDPCPPRoot) {
+    return createStringError("Could not locate DPCPP root directory");
+  }
+
+  constexpr std::array<llvm::StringLiteral, 8> SYCLDeviceWrapperLibs = {
+      "libsycl-crt",      "libsycl-complex",    "libsycl-complex-fp64",
+      "libsycl-cmath",    "libsycl-cmath-fp64", "libsycl-imf",
+      "libsycl-imf-fp64", "libsycl-imf-bf16"};
+
+  constexpr std::array<llvm::StringLiteral, 3> SYCLDeviceAnnotationLibs = {
+      "libsycl-itt-user-wrappers", "libsycl-itt-compiler-wrappers",
+      "libsycl-itt-stubs"};
+
+  LLVMContext &Context = Module.getContext();
+  auto Link = [&](ArrayRef<llvm::StringLiteral> LibNames) -> Error {
+    for (const auto &LibName : LibNames) {
+      std::string LibPath = (DPCPPRoot + "/lib/" + LibName + ".bc").str();
+
+      SMDiagnostic Diag;
+      std::unique_ptr<llvm::Module> Lib = parseIRFile(LibPath, Diag, Context);
+      if (!Lib) {
+        std::string DiagMsg;
+        raw_string_ostream SOS(DiagMsg);
+        Diag.print(/*ProgName=*/nullptr, SOS);
+        return createStringError(DiagMsg);
+      }
+
+      if (Linker::linkModules(Module, std::move(Lib), Linker::LinkOnlyNeeded)) {
+        // TODO: `linkModules` always prints errors to the console.
+        return createStringError("Unable to link device library: %s",
+                                 LibPath.c_str());
+      }
+    }
+
+    return Error::success();
+  };
+
+  if (auto Error = Link(SYCLDeviceWrapperLibs)) {
+    return Error;
+  }
+  if (DeviceInstrumentationEnabled) {
+    if (auto Error = Link(SYCLDeviceAnnotationLibs)) {
+      return Error;
+    }
+  }
+  return Error::success();
 }

sycl-jit/jit-compiler/lib/rtc/DeviceCompilation.h

@@ -23,6 +23,9 @@ llvm::Expected<std::unique_ptr<llvm::Module>>
 compileDeviceCode(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
                   View<const char *> UserArgs);
 
+llvm::Error linkDefaultDeviceLibraries(llvm::Module &Module,
+                                       View<const char *> UserArgs);
+
 } // namespace jit_compiler
 
 #endif // SYCL_JIT_COMPILER_RTC_DEVICE_COMPILATION_H