[SYCL] PoC implementation of kernel compiler extension with libtooling and sycl-jit

sycl-jit/common/include/Kernel.h

@@ -349,6 +349,11 @@ struct SYCLKernelInfo {
       : Name{KernelName}, Args{NumArgs}, Attributes{}, NDR{}, BinaryInfo{} {}
 };
 
+struct IncludePair {
+  const char *Path;
+  const char *Contents;
+};
+
 } // namespace jit_compiler
 
 #endif // SYCL_FUSION_COMMON_KERNEL_H

sycl-jit/jit-compiler/CMakeLists.txt

@@ -7,6 +7,7 @@ add_llvm_library(sycl-jit
    lib/fusion/FusionHelper.cpp
    lib/fusion/JITContext.cpp
    lib/fusion/ModuleHelper.cpp
+   lib/rtc/DeviceCompilation.cpp
    lib/helper/ConfigHelper.cpp
 
    SHARED
@@ -29,6 +30,14 @@ add_llvm_library(sycl-jit
    TargetParser
    MC
    ${LLVM_TARGETS_TO_BUILD}
+
+   LINK_LIBS
+   clangBasic
+   clangDriver
+   clangFrontend
+   clangCodeGen
+   clangTooling
+   clangSerialization
 )
 
 target_compile_options(sycl-jit PRIVATE ${SYCL_JIT_WARNING_FLAGS})
@@ -40,6 +49,8 @@ target_include_directories(sycl-jit
   SYSTEM PRIVATE
   ${LLVM_MAIN_INCLUDE_DIR}
   ${LLVM_SPIRV_INCLUDE_DIRS}
+  ${CMAKE_SOURCE_DIR}/../clang/include
+  ${CMAKE_BINARY_DIR}/tools/clang/include
 )
 target_include_directories(sycl-jit
   PUBLIC

sycl-jit/jit-compiler/include/KernelFusion.h

@@ -66,6 +66,9 @@ JITResult materializeSpecConstants(const char *KernelName,
                                    jit_compiler::SYCLKernelBinaryInfo &BinInfo,
                                    View<unsigned char> SpecConstBlob);
 
+JITResult compileSYCL(const char *SYCLSource, View<IncludePair> IncludePairs,
+                      View<const char *> UserArgs, const char *DPCPPRoot);
+
 /// Clear all previously set options.
 void resetJITConfiguration();
 

sycl-jit/jit-compiler/ld-version-script.txt

@@ -3,6 +3,7 @@
     /* Export the library entry points */
     fuseKernels;
     materializeSpecConstants;
+    compileSYCL;
     resetJITConfiguration;
     addToJITConfiguration;
 

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

@@ -14,6 +14,7 @@
 #include "fusion/FusionPipeline.h"
 #include "helper/ConfigHelper.h"
 #include "helper/ErrorHandling.h"
+#include "rtc/DeviceCompilation.h"
 #include "translation/KernelTranslation.h"
 #include "translation/SPIRVLLVMTranslation.h"
 #include <llvm/Support/Error.h>
@@ -235,6 +236,31 @@ extern "C" JITResult fuseKernels(View<SYCLKernelInfo> KernelInformation,
   return JITResult{FusedKernelInfo};
 }
 
+extern "C" JITResult compileSYCL(const char *SYCLSource,
+                                 View<IncludePair> IncludePairs,
+                                 View<const char *> UserArgs,
+                                 const char *DPCPPRoot) {
+  std::unique_ptr<llvm::Module> Module =
+      compileDeviceCode(SYCLSource, IncludePairs, UserArgs, DPCPPRoot);
+  if (!Module) {
+    return JITResult{"Device code compilation failed"};
+  }
+
+  SYCLKernelInfo Kernel;
+  auto Error = translation::KernelTranslator::translateKernel(
+      Kernel, *Module, JITContext::getInstance(), BinaryFormat::SPIRV);
+
+  auto *LLVMCtx = &Module->getContext();
+  Module.reset();
+  delete LLVMCtx;
+
+  if (Error) {
+    return errorToFusionResult(std::move(Error), "SPIR-V translation failed");
+  }
+
+  return JITResult{Kernel};
+}
+
 extern "C" void resetJITConfiguration() { ConfigHelper::reset(); }
 
 extern "C" void addToJITConfiguration(OptionStorage &&Opt) {

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

@@ -0,0 +1,111 @@
+//==---------------------- DeviceCompilation.cpp ---------------------------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "DeviceCompilation.h"
+
+#include <clang/Basic/Version.h>
+#include <clang/CodeGen/CodeGenAction.h>
+#include <clang/Driver/Compilation.h>
+#include <clang/Frontend/CompilerInstance.h>
+#include <clang/Tooling/CompilationDatabase.h>
+#include <clang/Tooling/Tooling.h>
+
+#include <llvm/IR/Module.h>
+
+namespace {
+using namespace clang;
+using namespace clang::tooling;
+using namespace clang::driver;
+
+struct GetLLVMModuleAction : public ToolAction {
+  // Code adapted from `FrontendActionFactory::runInvocation`.
+  bool runInvocation(std::shared_ptr<CompilerInvocation> Invocation,
+                     FileManager *Files,
+                     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
+                     DiagnosticConsumer *DiagConsumer) override {
+    assert(!Module && "Action should only be invoked on a single file");
+
+    // Create a compiler instance to handle the actual work.
+    CompilerInstance Compiler(std::move(PCHContainerOps));
+    Compiler.setInvocation(std::move(Invocation));
+    Compiler.setFileManager(Files);
+
+    // Create the compiler's actual diagnostics engine.
+    Compiler.createDiagnostics(DiagConsumer, /*ShouldOwnClient=*/false);
+    if (!Compiler.hasDiagnostics()) {
+      return false;
+    }
+
+    Compiler.createSourceManager(*Files);
+
+    // Ignore `Compiler.getFrontendOpts().ProgramAction` (would be `EmitBC`) and
+    // create/execute an `EmitLLVMOnlyAction` (= codegen to LLVM module without
+    // emitting anything) instead.
+    EmitLLVMOnlyAction ELOA;
+    const bool Success = Compiler.ExecuteAction(ELOA);
+    Files->clearStatCache();
+    if (!Success) {
+      return false;
+    }
+
+    // Take the module and its context to extend the objects' lifetime.
+    Module = ELOA.takeModule();
+    ELOA.takeLLVMContext();
+
+    return true;
+  }
+
+  std::unique_ptr<llvm::Module> Module;
+};
+
+} // anonymous namespace
+
+std::unique_ptr<llvm::Module> jit_compiler::compileDeviceCode(
+    const char *SYCLSource, View<IncludePair> IncludePairs,
+    View<const char *> UserArgs, const char *DPCPPRoot) {
+
+  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};
+
+  constexpr auto SourcePath = "rtc.cpp";
+  clang::tooling::ClangTool Tool{DB, {SourcePath}};
+
+  // Set up in-memory filesystem.
+  Tool.mapVirtualFile(SourcePath, SYCLSource);
+  for (const auto &IP : IncludePairs) {
+    Tool.mapVirtualFile(IP.Path, IP.Contents);
+  }
+
+  // Reset argument adjusters to drop the `-fsyntax-only` flag which is added by
+  // default by this API.
+  Tool.clearArgumentsAdjusters();
+  // Then, modify argv[0] and set the resource directory so that the driver
+  // picks up the correct SYCL environment.
+  Tool.appendArgumentsAdjuster(
+      [&DPCPPRoot](const CommandLineArguments &Args,
+                   StringRef Filename) -> CommandLineArguments {
+        (void)Filename;
+        CommandLineArguments NewArgs = Args;
+        NewArgs[0] = (Twine(DPCPPRoot) + "/bin/clang++").str();
+        NewArgs.push_back((Twine("-resource-dir=") + DPCPPRoot + "/lib/clang/" +
+                           Twine(CLANG_VERSION_MAJOR))
+                              .str());
+        return NewArgs;
+      });
+
+  GetLLVMModuleAction Action;
+  if (!Tool.run(&Action)) {
+    return std::move(Action.Module);
+  }
+
+  return {};
+}

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

@@ -0,0 +1,30 @@
+//==---- DeviceCompilation.h - Compile SYCL device code with libtooling ----==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SYCL_JIT_COMPILER_RTC_DEVICE_COMPILATION_H
+#define SYCL_JIT_COMPILER_RTC_DEVICE_COMPILATION_H
+
+#include "Kernel.h"
+#include "View.h"
+
+#include <memory>
+
+namespace llvm {
+class Module;
+} // namespace llvm
+
+namespace jit_compiler {
+
+std::unique_ptr<llvm::Module> compileDeviceCode(const char *SYCLSource,
+                                                View<IncludePair> IncludePairs,
+                                                View<const char *> UserArgs,
+                                                const char *DPCPPRoot);
+
+} // namespace jit_compiler
+
+#endif // SYCL_JIT_COMPILER_RTC_DEVICE_COMPILATION_H

sycl-jit/jit-compiler/lib/translation/SPIRVLLVMTranslation.cpp

@@ -41,6 +41,9 @@ SPIRV::TranslatorOpts &SPIRVLLVMTranslator::translatorOpts() {
     // there's currently no obvious way to iterate the
     // array of extensions in KernelInfo.
     TransOpt.enableAllExtensions();
+    // TODO: Remove this workaround.
+    TransOpt.setAllowedToUseExtension(
+        SPIRV::ExtensionID::SPV_KHR_untyped_pointers, false);
     TransOpt.setDesiredBIsRepresentation(
         SPIRV::BIsRepresentation::SPIRVFriendlyIR);
     // TODO: We need to take care of specialization constants, either by

sycl/source/detail/jit_compiler.cpp

@@ -18,6 +18,9 @@
 #include <sycl/detail/ur.hpp>
 #include <sycl/kernel_bundle.hpp>
 
+#include <dlfcn.h>
+#include <link.h>
+
 namespace sycl {
 inline namespace _V1 {
 namespace detail {
@@ -74,6 +77,31 @@ jit_compiler::jit_compiler() {
       return false;
     }
 
+    this->CompileSYCLHandle = reinterpret_cast<CompileSYCLFuncT>(
+        sycl::detail::ur::getOsLibraryFuncAddress(LibraryPtr, "compileSYCL"));
+    if (!this->CompileSYCLHandle) {
+      printPerformanceWarning(
+          "Cannot resolve JIT library function entry point");
+      return false;
+    }
+
+    // TODO: Move this query to a more appropriate location (e.g. add
+    // `sycl::detail::ur::getOsLibraryPath`), and handle non-POSIX OSs. For now,
+    // it should be fine here because the JIT is not built on Windows.
+    link_map *Map = nullptr;
+    if (dlinfo(LibraryPtr, RTLD_DI_LINKMAP, &Map) == 0) {
+      std::string LoadedLibraryPath = Map->l_name;
+      std::string JITLibraryPathSuffix = "/lib/" + JITLibraryName;
+      auto Pos = LoadedLibraryPath.rfind(JITLibraryPathSuffix);
+      if (Pos != std::string::npos) {
+        this->DPCPPRoot = LoadedLibraryPath.substr(0, Pos);
+      }
+    }
+    if (this->DPCPPRoot.empty()) {
+      printPerformanceWarning("Cannot determine JIT library location");
+      return false;
+    }
+
     return true;
   };
   Available = checkJITLibrary();
@@ -1143,6 +1171,45 @@ std::vector<uint8_t> jit_compiler::encodeReqdWorkGroupSize(
   return Encoded;
 }
 
+std::vector<uint8_t> jit_compiler::compileSYCL(
+    const std::string &SYCLSource,
+    const std::vector<std::pair<std::string, std::string>> &IncludePairs,
+    const std::vector<std::string> &UserArgs, std::string *LogPtr,
+    const std::vector<std::string> &RegisteredKernelNames) {
+
+  // TODO: Handle situation.
+  assert(RegisteredKernelNames.empty() &&
+         "Instantiation of kernel templates NYI");
+
+  std::vector<::jit_compiler::IncludePair> IncludePairsView;
+  IncludePairsView.reserve(IncludePairs.size());
+  std::transform(IncludePairs.begin(), IncludePairs.end(),
+                 std::back_inserter(IncludePairsView), [](const auto &Pair) {
+                   return ::jit_compiler::IncludePair{Pair.first.c_str(),
+                                                      Pair.second.c_str()};
+                 });
+  std::vector<const char *> UserArgsView;
+  UserArgsView.reserve(UserArgs.size());
+  std::transform(UserArgs.begin(), UserArgs.end(),
+                 std::back_inserter(UserArgsView),
+                 [](const auto &Arg) { return Arg.c_str(); });
+
+  auto Result = CompileSYCLHandle(SYCLSource.c_str(), IncludePairsView,
+                                  UserArgsView, DPCPPRoot.c_str());
+
+  if (Result.failed()) {
+    throw sycl::exception(sycl::errc::build, Result.getErrorMessage());
+  }
+
+  // TODO: We currently don't have a meaningful build log.
+  (void)LogPtr;
+
+  const auto &BI = Result.getKernelInfo().BinaryInfo;
+  assert(BI.Format == ::jit_compiler::BinaryFormat::SPIRV);
+  std::vector<uint8_t> SPV(BI.BinaryStart, BI.BinaryStart + BI.BinarySize);
+  return SPV;
+}
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl

sycl/source/detail/jit_compiler.hpp

@@ -44,6 +44,12 @@ class jit_compiler {
                            const std::string &KernelName,
                            const std::vector<unsigned char> &SpecConstBlob);
 
+  std::vector<uint8_t> compileSYCL(
+      const std::string &SYCLSource,
+      const std::vector<std::pair<std::string, std::string>> &IncludePairs,
+      const std::vector<std::string> &UserArgs, std::string *LogPtr,
+      const std::vector<std::string> &RegisteredKernelNames);
+
   bool isAvailable() { return Available; }
 
   static jit_compiler &get_instance() {
@@ -80,12 +86,15 @@ class jit_compiler {
   using FuseKernelsFuncT = decltype(::jit_compiler::fuseKernels) *;
   using MaterializeSpecConstFuncT =
       decltype(::jit_compiler::materializeSpecConstants) *;
+  using CompileSYCLFuncT = decltype(::jit_compiler::compileSYCL) *;
   using ResetConfigFuncT = decltype(::jit_compiler::resetJITConfiguration) *;
   using AddToConfigFuncT = decltype(::jit_compiler::addToJITConfiguration) *;
   FuseKernelsFuncT FuseKernelsHandle = nullptr;
   MaterializeSpecConstFuncT MaterializeSpecConstHandle = nullptr;
+  CompileSYCLFuncT CompileSYCLHandle = nullptr;
   ResetConfigFuncT ResetConfigHandle = nullptr;
   AddToConfigFuncT AddToConfigHandle = nullptr;
+  std::string DPCPPRoot;
 #endif // SYCL_EXT_JIT_ENABLE
 };
 

sycl/source/detail/kernel_bundle_impl.hpp

@@ -446,6 +446,12 @@ class kernel_bundle_impl {
                                              BuildOptions, LogPtr,
                                              RegisteredKernelNames);
       }
+      if (Language == syclex::source_language::sycljit) {
+        const auto &SourceStr = std::get<std::string>(this->Source);
+        return syclex::detail::SYCLJIT_to_SPIRV(SourceStr, IncludePairs,
+                                                BuildOptions, LogPtr,
+                                                RegisteredKernelNames);
+      }
       throw sycl::exception(
           make_error_code(errc::invalid),
           "OpenCL C and SPIR-V are the only supported languages at this time");