[SYCL] RTC support for GPU targets

This patch extends RTC support to GPU (AMD and Nvidia) targets.

Additionally:
* reinstate __SYCL_PROGRAM_METADATA_TAG_NEED_FINALIZATION tag,
* split sycl.cpp RTC file to exclude IMF from the body of the main test.

Author: Jakub Chlanda

sycl-jit/jit-compiler/CMakeLists.txt

@@ -60,6 +60,7 @@ target_include_directories(sycl-jit
   ${LLVM_MAIN_INCLUDE_DIR}
   ${LLVM_SPIRV_INCLUDE_DIRS}
   ${LLVM_EXTERNAL_CLANG_SOURCE_DIR}/include
+  ${LLVM_EXTERNAL_CLANG_SOURCE_DIR}/lib
   ${CMAKE_BINARY_DIR}/tools/clang/include
 )
 target_include_directories(sycl-jit

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

@@ -176,10 +176,11 @@ class RTCResult {
 
 /// Calculates a BLAKE3 hash of the pre-processed source string described by
 /// \p SourceFile (considering any additional \p IncludeFiles) and the
-/// concatenation of the \p UserArgs.
+/// concatenation of the \p UserArgs, for a given \p Format.
 JIT_EXPORT_SYMBOL RTCHashResult calculateHash(InMemoryFile SourceFile,
                                               View<InMemoryFile> IncludeFiles,
-                                              View<const char *> UserArgs);
+                                              View<const char *> UserArgs,
+                                              BinaryFormat Format);
 
 /// Compiles, links against device libraries, and finalizes the device code in
 /// the source string described by \p SourceFile, considering any additional \p
@@ -191,10 +192,14 @@ JIT_EXPORT_SYMBOL RTCHashResult calculateHash(InMemoryFile SourceFile,
 ///
 /// If \p SaveIR is true and \p CachedIR is empty, the LLVM module obtained from
 /// the frontend invocation is wrapped in bitcode format in the result object.
+///
+/// \p BinaryFormat describes the desired format of the compilation - which
+/// corresponds to the backend that is being targeted.
 JIT_EXPORT_SYMBOL RTCResult compileSYCL(InMemoryFile SourceFile,
                                         View<InMemoryFile> IncludeFiles,
                                         View<const char *> UserArgs,
-                                        View<char> CachedIR, bool SaveIR);
+                                        View<char> CachedIR, bool SaveIR,
+                                        BinaryFormat Format);
 
 /// Requests that the JIT binary referenced by \p Address is deleted from the
 /// `JITContext`.

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

@@ -8,6 +8,8 @@
 
 #include "DeviceCompilation.h"
 #include "ESIMD.h"
+#include "JITBinaryInfo.h"
+#include "translation/Translation.h"
 
 #include <clang/Basic/DiagnosticDriver.h>
 #include <clang/Basic/Version.h>
@@ -22,6 +24,15 @@
 #include <clang/Frontend/Utils.h>
 #include <clang/Tooling/CompilationDatabase.h>
 #include <clang/Tooling/Tooling.h>
+#if defined(JIT_SUPPORT_PTX) || defined(JIT_SUPPORT_AMDGCN)
+#include <clang/Driver/Driver.h>
+#endif
+#ifdef JIT_SUPPORT_PTX
+#include <Driver/ToolChains/Cuda.h>
+#include <Driver/ToolChains/LazyDetector.h>
+#elif JIT_SUPPORT_AMDGCN
+#include <Driver/ToolChains/AMDGPU.h>
+#endif
 
 #include <llvm/IR/DiagnosticInfo.h>
 #include <llvm/IR/DiagnosticPrinter.h>
@@ -178,7 +189,8 @@ class RTCToolActionBase : public ToolAction {
     assert(!hasExecuted() && "Action should only be invoked on a single file");
 
     // Create a compiler instance to handle the actual work.
-    CompilerInstance Compiler(std::move(Invocation), std::move(PCHContainerOps));
+    CompilerInstance Compiler(std::move(Invocation),
+                              std::move(PCHContainerOps));
     Compiler.setFileManager(Files);
     // Suppress summary with number of warnings and errors being printed to
     // stdout.
@@ -361,10 +373,24 @@ static void setupTool(ClangTool &Tool, const std::string &DPCPPRoot,
       });
 }
 
-Expected<std::string>
-jit_compiler::calculateHash(InMemoryFile SourceFile,
-                            View<InMemoryFile> IncludeFiles,
-                            const InputArgList &UserArgList) {
+static void setGPUTarget(BinaryFormat Format,
+                         SmallVector<std::string> &CommandLine) {
+  auto [CPU, _] = Translator::getTargetCPUAndFeatureAttrs(nullptr, "", Format);
+  CommandLine.push_back("-fsycl");
+  if (Format == BinaryFormat::PTX) {
+    CommandLine.push_back("-fsycl-targets=nvptx64-nvidia-cuda");
+    CommandLine.push_back("-Xsycl-target-backend");
+    CommandLine.push_back("--cuda-gpu-arch=" + CPU);
+  } else if (Format == BinaryFormat::AMDGCN) {
+    CommandLine.push_back("-fsycl-targets=amdgcn-amd-amdhsa");
+    CommandLine.push_back("-Xsycl-target-backend=amdgcn-amd-amdhsa");
+    CommandLine.push_back("--offload-arch=" + CPU);
+  }
+}
+
+Expected<std::string> jit_compiler::calculateHash(
+    InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
+    const InputArgList &UserArgList, BinaryFormat Format) {
   TimeTraceScope TTS{"calculateHash"};
 
   const std::string &DPCPPRoot = getDPCPPRoot();
@@ -373,6 +399,9 @@ jit_compiler::calculateHash(InMemoryFile SourceFile,
   }
 
   SmallVector<std::string> CommandLine;
+  if (Format == BinaryFormat::PTX || Format == BinaryFormat::AMDGCN) {
+    setGPUTarget(Format, CommandLine);
+  }
   adjustArgs(UserArgList, DPCPPRoot, CommandLine);
 
   FixedCompilationDatabase DB{".", CommandLine};
@@ -399,11 +428,10 @@ jit_compiler::calculateHash(InMemoryFile SourceFile,
   return createStringError("Calculating source hash failed");
 }
 
-Expected<ModuleUPtr>
-jit_compiler::compileDeviceCode(InMemoryFile SourceFile,
-                                View<InMemoryFile> IncludeFiles,
-                                const InputArgList &UserArgList,
-                                std::string &BuildLog, LLVMContext &Context) {
+Expected<ModuleUPtr> jit_compiler::compileDeviceCode(
+    InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
+    const InputArgList &UserArgList, std::string &BuildLog,
+    LLVMContext &Context, BinaryFormat Format) {
   TimeTraceScope TTS{"compileDeviceCode"};
 
   const std::string &DPCPPRoot = getDPCPPRoot();
@@ -412,6 +440,9 @@ jit_compiler::compileDeviceCode(InMemoryFile SourceFile,
   }
 
   SmallVector<std::string> CommandLine;
+  if (Format == BinaryFormat::PTX || Format == BinaryFormat::AMDGCN) {
+    setGPUTarget(Format, CommandLine);
+  }
   adjustArgs(UserArgList, DPCPPRoot, CommandLine);
 
   FixedCompilationDatabase DB{".", CommandLine};
@@ -430,12 +461,22 @@ jit_compiler::compileDeviceCode(InMemoryFile SourceFile,
   return createStringError(BuildLog);
 }
 
-// This function is a simplified copy of the device library selection process in
-// `clang::driver::tools::SYCL::getDeviceLibraries`, assuming a SPIR-V target
-// (no AoT, no third-party GPUs, no native CPU). Keep in sync!
+// This function is a simplified copy of the device library selection process
+// in `clang::driver::tools::SYCL::getDeviceLibraries`, assuming a SPIR-V, or
+// GPU targets (no AoT, no native CPU). Keep in sync!
 static bool getDeviceLibraries(const ArgList &Args,
                                SmallVectorImpl<std::string> &LibraryList,
-                               DiagnosticsEngine &Diags) {
+                               DiagnosticsEngine &Diags, BinaryFormat Format) {
+  // For CUDA/HIP we only need devicelib, early exit here.
+  if (Format == BinaryFormat::PTX) {
+    LibraryList.push_back(
+        Args.MakeArgString("devicelib-nvptx64-nvidia-cuda.bc"));
+    return false;
+  } else if (Format == BinaryFormat::AMDGCN) {
+    LibraryList.push_back(Args.MakeArgString("devicelib-amdgcn-amd-amdhsa.bc"));
+    return false;
+  }
+
   struct DeviceLibOptInfo {
     StringRef DeviceLibName;
     StringRef DeviceLibOption;
@@ -540,7 +581,8 @@ static Expected<ModuleUPtr> loadBitcodeLibrary(StringRef LibPath,
 
 Error jit_compiler::linkDeviceLibraries(llvm::Module &Module,
                                         const InputArgList &UserArgList,
-                                        std::string &BuildLog) {
+                                        std::string &BuildLog,
+                                        BinaryFormat Format) {
   TimeTraceScope TTS{"linkDeviceLibraries"};
 
   const std::string &DPCPPRoot = getDPCPPRoot();
@@ -555,11 +597,29 @@ Error jit_compiler::linkDeviceLibraries(llvm::Module &Module,
                           /* ShouldOwnClient=*/false);
 
   SmallVector<std::string> LibNames;
-  bool FoundUnknownLib = getDeviceLibraries(UserArgList, LibNames, Diags);
+  const bool FoundUnknownLib =
+      getDeviceLibraries(UserArgList, LibNames, Diags, Format);
   if (FoundUnknownLib) {
     return createStringError("Could not determine list of device libraries: %s",
                              BuildLog.c_str());
   }
+  const bool IsGPUTarget =
+      Format == BinaryFormat::PTX || Format == BinaryFormat::AMDGCN;
+  if (IsGPUTarget) {
+    // Based on the OS and the format decide on the version of libspirv.
+    // NOTE: this will be problematic if cross-compiling between OSes.
+    std::string Libclc{"clc/"};
+    Libclc.append(
+#ifdef _WIN32
+        "remangled-l32-signed_char.libspirv-"
+#else
+        "remangled-l64-signed_char.libspirv-"
+#endif
+    );
+    Libclc.append(Format == BinaryFormat::PTX ? "nvptx64-nvidia-cuda.bc"
+                                              : "amdgcn-amd-amdhsa.bc");
+    LibNames.push_back(Libclc);
+  }
 
   LLVMContext &Context = Module.getContext();
   for (const std::string &LibName : LibNames) {
@@ -577,6 +637,57 @@ Error jit_compiler::linkDeviceLibraries(llvm::Module &Module,
     }
   }
 
+  // For GPU targets we need to link against vendor provided libdevice.
+  if (IsGPUTarget) {
+    Triple T{Module.getTargetTriple()};
+    Driver D{(Twine(DPCPPRoot) + "/bin/clang++").str(), T.getTriple(), Diags};
+    auto [CPU, _] =
+        Translator::getTargetCPUAndFeatureAttrs(&Module, "", Format);
+    // Helper lambda to link modules.
+    auto LinkInLib = [&](const StringRef LibDevice) -> Error {
+      ModuleUPtr LibDeviceModule;
+      if (auto Error = loadBitcodeLibrary(LibDevice, Context)
+                           .moveInto(LibDeviceModule)) {
+        return Error;
+      }
+      if (Linker::linkModules(Module, std::move(LibDeviceModule),
+                              Linker::LinkOnlyNeeded)) {
+        return createStringError("Unable to link libdevice: %s",
+                                 BuildLog.c_str());
+      }
+      return Error::success();
+    };
+    SmallVector<std::string, 12> LibDeviceFiles;
+#ifdef JIT_SUPPORT_PTX
+    // For NVPTX we can get away with CudaInstallationDetector.
+    LazyDetector<CudaInstallationDetector> CudaInstallation{D, T, UserArgList};
+    auto LibDevice = CudaInstallation->getLibDeviceFile(CPU);
+    if (LibDevice.empty()) {
+      return createStringError("Unable to find Cuda libdevice");
+    }
+    LibDeviceFiles.push_back(LibDevice);
+#elif JIT_SUPPORT_AMDGCN
+    // AMDGPU requires entire toolchain in order to provide all common bitcode
+    // libraries.
+    clang::driver::toolchains::ROCMToolChain TC(D, T, UserArgList);
+    auto CommonDeviceLibs = TC.getCommonDeviceLibNames(
+        UserArgList, CPU, Action::OffloadKind::OFK_SYCL, false);
+    if (CommonDeviceLibs.empty()) {
+      return createStringError("Unable to find ROCm common device libraries");
+    }
+    for (auto &Lib : CommonDeviceLibs) {
+      LibDeviceFiles.push_back(Lib.Path);
+    }
+#endif
+    for (auto &LibDeviceFile : LibDeviceFiles) {
+      auto Res = LinkInLib(LibDeviceFile);
+      // llvm::Error converts to false on success.
+      if (Res) {
+        return Res;
+      }
+    }
+  }
+
   return Error::success();
 }
 

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

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "JITBinaryInfo.h"
 #include "RTC.h"
 
 #include <llvm/ADT/SmallVector.h>
@@ -24,16 +25,17 @@ using ModuleUPtr = std::unique_ptr<llvm::Module>;
 
 llvm::Expected<std::string>
 calculateHash(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
-              const llvm::opt::InputArgList &UserArgList);
+              const llvm::opt::InputArgList &UserArgList, BinaryFormat Format);
 
 llvm::Expected<ModuleUPtr>
 compileDeviceCode(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
                   const llvm::opt::InputArgList &UserArgList,
-                  std::string &BuildLog, llvm::LLVMContext &Context);
+                  std::string &BuildLog, llvm::LLVMContext &Context,
+                  BinaryFormat Format);
 
 llvm::Error linkDeviceLibraries(llvm::Module &Module,
                                 const llvm::opt::InputArgList &UserArgList,
-                                std::string &BuildLog);
+                                std::string &BuildLog, BinaryFormat Format);
 
 using PostLinkResult = std::pair<RTCBundleInfo, llvm::SmallVector<ModuleUPtr>>;
 llvm::Expected<PostLinkResult>

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

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "RTC.h"
+#include "JITBinaryInfo.h"
 #include "helper/ErrorHelper.h"
 #include "rtc/DeviceCompilation.h"
 #include "translation/SPIRVLLVMTranslation.h"
@@ -26,7 +27,8 @@ using namespace jit_compiler;
 
 JIT_EXPORT_SYMBOL RTCHashResult calculateHash(InMemoryFile SourceFile,
                                               View<InMemoryFile> IncludeFiles,
-                                              View<const char *> UserArgs) {
+                                              View<const char *> UserArgs,
+                                              BinaryFormat Format) {
   llvm::opt::InputArgList UserArgList;
   if (auto Error = parseUserArgs(UserArgs).moveInto(UserArgList)) {
     return errorTo<RTCHashResult>(std::move(Error),
@@ -36,8 +38,8 @@ JIT_EXPORT_SYMBOL RTCHashResult calculateHash(InMemoryFile SourceFile,
 
   auto Start = std::chrono::high_resolution_clock::now();
   std::string Hash;
-  if (auto Error =
-          calculateHash(SourceFile, IncludeFiles, UserArgList).moveInto(Hash)) {
+  if (auto Error = calculateHash(SourceFile, IncludeFiles, UserArgList, Format)
+                       .moveInto(Hash)) {
     return errorTo<RTCHashResult>(std::move(Error), "Hashing failed",
                                   /*IsHash=*/false);
   }
@@ -55,7 +57,8 @@ JIT_EXPORT_SYMBOL RTCHashResult calculateHash(InMemoryFile SourceFile,
 JIT_EXPORT_SYMBOL RTCResult compileSYCL(InMemoryFile SourceFile,
                                         View<InMemoryFile> IncludeFiles,
                                         View<const char *> UserArgs,
-                                        View<char> CachedIR, bool SaveIR) {
+                                        View<char> CachedIR, bool SaveIR,
+                                        BinaryFormat Format) {
   llvm::LLVMContext Context;
   std::string BuildLog;
   configureDiagnostics(Context, BuildLog);
@@ -104,7 +107,7 @@ JIT_EXPORT_SYMBOL RTCResult compileSYCL(InMemoryFile SourceFile,
   bool FromSource = !Module;
   if (FromSource) {
     if (auto Error = compileDeviceCode(SourceFile, IncludeFiles, UserArgList,
-                                       BuildLog, Context)
+                                       BuildLog, Context, Format)
                          .moveInto(Module)) {
       return errorTo<RTCResult>(std::move(Error), "Device compilation failed");
     }
@@ -118,7 +121,8 @@ JIT_EXPORT_SYMBOL RTCResult compileSYCL(InMemoryFile SourceFile,
     IR = RTCDeviceCodeIR{BCString.data(), BCString.data() + BCString.size()};
   }
 
-  if (auto Error = linkDeviceLibraries(*Module, UserArgList, BuildLog)) {
+  if (auto Error =
+          linkDeviceLibraries(*Module, UserArgList, BuildLog, Format)) {
     return errorTo<RTCResult>(std::move(Error), "Device linking failed");
   }
 
@@ -131,9 +135,9 @@ JIT_EXPORT_SYMBOL RTCResult compileSYCL(InMemoryFile SourceFile,
 
   for (auto [DevImgInfo, Module] :
        llvm::zip_equal(BundleInfo.DevImgInfos, Modules)) {
-    if (auto Error = Translator::translate(*Module, JITContext::getInstance(),
-                                           BinaryFormat::SPIRV)
-                         .moveInto(DevImgInfo.BinaryInfo)) {
+    if (auto Error =
+            Translator::translate(*Module, JITContext::getInstance(), Format)
+                .moveInto(DevImgInfo.BinaryInfo)) {
       return errorTo<RTCResult>(std::move(Error), "SPIR-V translation failed");
     }
   }