[SYCL][RTC] Use program manager (#16316)

The idea is to assign a per-compilation prefix (e.g., `rtc_42$`) to all
offload entries in the `sycl_device_binaries` datastructure before
feeding them into `ProgramManager::addImages`, resulting in unique
`kernel_id`s as far as the PM is concerned. When querying the PM for the
device images to construct the kernel bundle, I look for kernel IDs
starting with the current prefix, which should reliably return only the
device images corresponding to current compilation request.

Note that the actual kernel names don't change, i.e. `__sycl_kernel_foo`
keeps that name even though the PM might know it as
`rtc_42$__sycl_kernel_foo`. The prefix is stored inside the bundle, and
prepended to the requested kernel name in the
`ext_onapi_[has|get]_kernel(string)` methods.

Kernel objects are also obtained via the program manager. Compared to
creating the UR kernel from the selected device image's UR program
directly, this approach ensures eliminated arguments are handled
correctly. Hence, I was able to drop the previously mandatory
`-fno-sycl-dead-args-optimization` from the pipeline.

The compilation pipeline and extended kernel bundle now support multiple
device images. To test this, I added support for the
`-fsycl-device-code-split=` option, and apply it to one of the
compilations in the E2E test.

The persistent cache is circumvented for now for the `sycl_jit` language
(lack of suitable on-disk format), but should be brought back in the
future.

---------

Signed-off-by: Julian Oppermann <[email protected]>
Co-authored-by: Lukas Sommer <[email protected]>

Author: jopperm

sycl-jit/common/include/Kernel.h

@@ -359,7 +359,7 @@ struct InMemoryFile {
   const char *Contents;
 };
 
-using RTCBundleBinaryInfo = SYCLKernelBinaryInfo;
+using RTCDevImgBinaryInfo = SYCLKernelBinaryInfo;
 using FrozenSymbolTable = DynArray<sycl::detail::string>;
 
 // Note: `FrozenPropertyValue` and `FrozenPropertySet` constructors take
@@ -399,16 +399,18 @@ struct FrozenPropertySet {
 
 using FrozenPropertyRegistry = DynArray<FrozenPropertySet>;
 
-struct RTCBundleInfo {
-  RTCBundleBinaryInfo BinaryInfo;
+struct RTCDevImgInfo {
+  RTCDevImgBinaryInfo BinaryInfo;
   FrozenSymbolTable SymbolTable;
   FrozenPropertyRegistry Properties;
 
-  RTCBundleInfo() = default;
-  RTCBundleInfo(RTCBundleInfo &&) = default;
-  RTCBundleInfo &operator=(RTCBundleInfo &&) = default;
+  RTCDevImgInfo() = default;
+  RTCDevImgInfo(RTCDevImgInfo &&) = default;
+  RTCDevImgInfo &operator=(RTCDevImgInfo &&) = default;
 };
 
+using RTCBundleInfo = DynArray<RTCDevImgInfo>;
+
 } // namespace jit_compiler
 
 #endif // SYCL_FUSION_COMMON_KERNEL_H

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

@@ -266,17 +266,18 @@ compileSYCL(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
     return errorTo<RTCResult>(PostLinkResultOrError.takeError(),
                               "Post-link phase failed");
   }
-  RTCBundleInfo BundleInfo;
-  std::tie(BundleInfo, Module) = std::move(*PostLinkResultOrError);
-
-  auto BinaryInfoOrError =
-      translation::KernelTranslator::translateBundleToSPIRV(
-          *Module, JITContext::getInstance());
-  if (!BinaryInfoOrError) {
-    return errorTo<RTCResult>(BinaryInfoOrError.takeError(),
-                              "SPIR-V translation failed");
+  auto [BundleInfo, Modules] = std::move(*PostLinkResultOrError);
+
+  for (auto [DevImgInfo, Module] : llvm::zip_equal(BundleInfo, Modules)) {
+    auto BinaryInfoOrError =
+        translation::KernelTranslator::translateDevImgToSPIRV(
+            *Module, JITContext::getInstance());
+    if (!BinaryInfoOrError) {
+      return errorTo<RTCResult>(BinaryInfoOrError.takeError(),
+                                "SPIR-V translation failed");
+    }
+    DevImgInfo.BinaryInfo = std::move(*BinaryInfoOrError);
   }
-  BundleInfo.BinaryInfo = std::move(*BinaryInfoOrError);
 
   return RTCResult{std::move(BundleInfo), BuildLog.c_str()};
 }

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

@@ -233,8 +233,6 @@ Expected<std::unique_ptr<llvm::Module>> jit_compiler::compileDeviceCode(
   DerivedArgList DAL{UserArgList};
   const auto &OptTable = getDriverOptTable();
   DAL.AddFlagArg(nullptr, OptTable.getOption(OPT_fsycl_device_only));
-  DAL.AddFlagArg(nullptr,
-                 OptTable.getOption(OPT_fno_sycl_dead_args_optimization));
   DAL.AddJoinedArg(
       nullptr, OptTable.getOption(OPT_resource_dir_EQ),
       (DPCPPRoot + "/lib/clang/" + Twine(CLANG_VERSION_MAJOR)).str());
@@ -436,15 +434,35 @@ template <class PassClass> static bool runModulePass(llvm::Module &M) {
   return !Res.areAllPreserved();
 }
 
-llvm::Expected<PostLinkResult> jit_compiler::performPostLink(
-    std::unique_ptr<llvm::Module> Module,
-    [[maybe_unused]] const llvm::opt::InputArgList &UserArgList) {
+static IRSplitMode getDeviceCodeSplitMode(const InputArgList &UserArgList) {
+  // This is the (combined) logic from
+  // `get[NonTriple|Triple]BasedSYCLPostLinkOpts` in
+  // `clang/lib/Driver/ToolChains/Clang.cpp`: Default is auto mode, but the user
+  // can override it by specifying the `-fsycl-device-code-split=` option. The
+  // no-argument variant `-fsycl-device-code-split` is ignored.
+  if (auto *Arg = UserArgList.getLastArg(OPT_fsycl_device_code_split_EQ)) {
+    StringRef ArgVal{Arg->getValue()};
+    if (ArgVal == "per_kernel") {
+      return SPLIT_PER_KERNEL;
+    }
+    if (ArgVal == "per_source") {
+      return SPLIT_PER_TU;
+    }
+    if (ArgVal == "off") {
+      return SPLIT_NONE;
+    }
+  }
+  return SPLIT_AUTO;
+}
+
+Expected<PostLinkResult>
+jit_compiler::performPostLink(std::unique_ptr<llvm::Module> Module,
+                              const InputArgList &UserArgList) {
   // This is a simplified version of `processInputModule` in
   // `llvm/tools/sycl-post-link.cpp`. Assertions/TODOs point to functionality
   // left out of the algorithm for now.
 
-  // TODO: SplitMode can be controlled by the user.
-  const auto SplitMode = SPLIT_NONE;
+  const auto SplitMode = getDeviceCodeSplitMode(UserArgList);
 
   // TODO: EmitOnlyKernelsAsEntryPoints is controlled by
   //       `shouldEmitOnlyKernelsAsEntryPoints` in
@@ -480,77 +498,87 @@ llvm::Expected<PostLinkResult> jit_compiler::performPostLink(
     return createStringError("`invoke_simd` calls detected");
   }
 
-  // TODO: Implement actual device code splitting. We're just using the splitter
-  //       to obtain additional information about the module for now.
-
   std::unique_ptr<ModuleSplitterBase> Splitter = getDeviceCodeSplitter(
       ModuleDesc{std::move(Module)}, SplitMode,
       /*IROutputOnly=*/false, EmitOnlyKernelsAsEntryPoints);
   assert(Splitter->hasMoreSplits());
-  if (Splitter->remainingSplits() > 1) {
-    return createStringError("Device code requires splitting");
-  }
 
   // TODO: Call `verifyNoCrossModuleDeviceGlobalUsage` if device globals shall
   //       be processed.
 
-  ModuleDesc MDesc = Splitter->nextSplit();
+  // TODO: This allocation assumes that there are no further splits required,
+  //       i.e. there are no mixed SYCL/ESIMD modules.
+  RTCBundleInfo BundleInfo{Splitter->remainingSplits()};
+  SmallVector<std::unique_ptr<llvm::Module>> Modules;
 
-  // TODO: Call `MDesc.fixupLinkageOfDirectInvokeSimdTargets()` when
-  //       `invoke_simd` is supported.
+  auto *DevImgInfoIt = BundleInfo.begin();
+  while (Splitter->hasMoreSplits()) {
+    assert(DevImgInfoIt != BundleInfo.end());
 
-  SmallVector<ModuleDesc, 2> ESIMDSplits =
-      splitByESIMD(std::move(MDesc), EmitOnlyKernelsAsEntryPoints);
-  assert(!ESIMDSplits.empty());
-  if (ESIMDSplits.size() > 1) {
-    return createStringError("Mixing SYCL and ESIMD code is unsupported");
-  }
-  MDesc = std::move(ESIMDSplits.front());
+    ModuleDesc MDesc = Splitter->nextSplit();
+    RTCDevImgInfo &DevImgInfo = *DevImgInfoIt++;
 
-  if (MDesc.isESIMD()) {
-    // `sycl-post-link` has a `-lower-esimd` option, but there's no clang driver
-    // option to influence it. Rather, the driver sets it unconditionally in the
-    // multi-file output mode, which we are mimicking here.
-    lowerEsimdConstructs(MDesc, PerformOpts);
-  }
+    // TODO: Call `MDesc.fixupLinkageOfDirectInvokeSimdTargets()` when
+    //       `invoke_simd` is supported.
 
-  MDesc.saveSplitInformationAsMetadata();
-
-  RTCBundleInfo BundleInfo;
-  BundleInfo.SymbolTable = FrozenSymbolTable{MDesc.entries().size()};
-  transform(MDesc.entries(), BundleInfo.SymbolTable.begin(),
-            [](Function *F) { return F->getName(); });
-
-  // TODO: Determine what is requested.
-  GlobalBinImageProps PropReq{
-      /*EmitKernelParamInfo=*/true, /*EmitProgramMetadata=*/true,
-      /*EmitExportedSymbols=*/true, /*EmitImportedSymbols=*/true,
-      /*DeviceGlobals=*/false};
-  PropertySetRegistry Properties =
-      computeModuleProperties(MDesc.getModule(), MDesc.entries(), PropReq);
-  // TODO: Manually add `compile_target` property as in
-  //       `saveModuleProperties`?
-  const auto &PropertySets = Properties.getPropSets();
-
-  BundleInfo.Properties = FrozenPropertyRegistry{PropertySets.size()};
-  for (auto &&[KV, FrozenPropSet] : zip(PropertySets, BundleInfo.Properties)) {
-    const auto &PropertySetName = KV.first;
-    const auto &PropertySet = KV.second;
-    FrozenPropSet =
-        FrozenPropertySet{PropertySetName.str(), PropertySet.size()};
-    for (auto &&[KV2, FrozenProp] : zip(PropertySet, FrozenPropSet.Values)) {
-      const auto &PropertyName = KV2.first;
-      const auto &PropertyValue = KV2.second;
-      FrozenProp = PropertyValue.getType() == PropertyValue::Type::UINT32
-                       ? FrozenPropertyValue{PropertyName.str(),
-                                             PropertyValue.asUint32()}
-                       : FrozenPropertyValue{
-                             PropertyName.str(), PropertyValue.asRawByteArray(),
-                             PropertyValue.getRawByteArraySize()};
+    SmallVector<ModuleDesc, 2> ESIMDSplits =
+        splitByESIMD(std::move(MDesc), EmitOnlyKernelsAsEntryPoints);
+    assert(!ESIMDSplits.empty());
+    if (ESIMDSplits.size() > 1) {
+      return createStringError("Mixing SYCL and ESIMD code is unsupported");
     }
-  };
+    MDesc = std::move(ESIMDSplits.front());
+
+    if (MDesc.isESIMD()) {
+      // `sycl-post-link` has a `-lower-esimd` option, but there's no clang
+      // driver option to influence it. Rather, the driver sets it
+      // unconditionally in the multi-file output mode, which we are mimicking
+      // here.
+      lowerEsimdConstructs(MDesc, PerformOpts);
+    }
+
+    MDesc.saveSplitInformationAsMetadata();
+
+    DevImgInfo.SymbolTable = FrozenSymbolTable{MDesc.entries().size()};
+    transform(MDesc.entries(), DevImgInfo.SymbolTable.begin(),
+              [](Function *F) { return F->getName(); });
+
+    // TODO: Determine what is requested.
+    GlobalBinImageProps PropReq{
+        /*EmitKernelParamInfo=*/true, /*EmitProgramMetadata=*/true,
+        /*EmitExportedSymbols=*/true, /*EmitImportedSymbols=*/true,
+        /*DeviceGlobals=*/false};
+    PropertySetRegistry Properties =
+        computeModuleProperties(MDesc.getModule(), MDesc.entries(), PropReq);
+    // TODO: Manually add `compile_target` property as in
+    //       `saveModuleProperties`?
+    const auto &PropertySets = Properties.getPropSets();
+
+    DevImgInfo.Properties = FrozenPropertyRegistry{PropertySets.size()};
+    for (auto [KV, FrozenPropSet] :
+         zip_equal(PropertySets, DevImgInfo.Properties)) {
+      const auto &PropertySetName = KV.first;
+      const auto &PropertySet = KV.second;
+      FrozenPropSet =
+          FrozenPropertySet{PropertySetName.str(), PropertySet.size()};
+      for (auto [KV2, FrozenProp] :
+           zip_equal(PropertySet, FrozenPropSet.Values)) {
+        const auto &PropertyName = KV2.first;
+        const auto &PropertyValue = KV2.second;
+        FrozenProp =
+            PropertyValue.getType() == PropertyValue::Type::UINT32
+                ? FrozenPropertyValue{PropertyName.str(),
+                                      PropertyValue.asUint32()}
+                : FrozenPropertyValue{PropertyName.str(),
+                                      PropertyValue.asRawByteArray(),
+                                      PropertyValue.getRawByteArraySize()};
+      }
+    };
+
+    Modules.push_back(MDesc.releaseModulePtr());
+  }
 
-  return PostLinkResult{std::move(BundleInfo), MDesc.releaseModulePtr()};
+  return PostLinkResult{std::move(BundleInfo), std::move(Modules)};
 }
 
 Expected<InputArgList>
@@ -607,21 +635,10 @@ jit_compiler::parseUserArgs(View<const char *> UserArgs) {
     }
   }
 
-  if (auto DCSMode = AL.getLastArgValue(OPT_fsycl_device_code_split_EQ, "none");
-      DCSMode != "none" && DCSMode != "auto") {
-    return createStringError("Device code splitting is not yet supported");
-  }
-
   if (!AL.hasFlag(OPT_fsycl_device_code_split_esimd,
                   OPT_fno_sycl_device_code_split_esimd, true)) {
     return createStringError("ESIMD device code split cannot be deactivated");
   }
 
-  if (AL.hasFlag(OPT_fsycl_dead_args_optimization,
-                 OPT_fno_sycl_dead_args_optimization, false)) {
-    return createStringError(
-        "Dead argument optimization must be disabled for runtime compilation");
-  }
-
   return std::move(AL);
 }

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

@@ -12,6 +12,7 @@
 #include "Kernel.h"
 #include "View.h"
 
+#include <llvm/ADT/SmallVector.h>
 #include <llvm/IR/Module.h>
 #include <llvm/Option/ArgList.h>
 #include <llvm/Support/Error.h>
@@ -30,7 +31,8 @@ llvm::Error linkDeviceLibraries(llvm::Module &Module,
                                 const llvm::opt::InputArgList &UserArgList,
                                 std::string &BuildLog);
 
-using PostLinkResult = std::pair<RTCBundleInfo, std::unique_ptr<llvm::Module>>;
+using PostLinkResult =
+    std::pair<RTCBundleInfo, llvm::SmallVector<std::unique_ptr<llvm::Module>>>;
 llvm::Expected<PostLinkResult>
 performPostLink(std::unique_ptr<llvm::Module> Module,
                 const llvm::opt::InputArgList &UserArgList);

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

@@ -222,18 +222,18 @@ llvm::Error KernelTranslator::translateKernel(SYCLKernelInfo &Kernel,
   return Error::success();
 }
 
-llvm::Expected<RTCBundleBinaryInfo>
-KernelTranslator::translateBundleToSPIRV(llvm::Module &Mod,
+llvm::Expected<RTCDevImgBinaryInfo>
+KernelTranslator::translateDevImgToSPIRV(llvm::Module &Mod,
                                          JITContext &JITCtx) {
   llvm::Expected<KernelBinary *> BinaryOrError = translateToSPIRV(Mod, JITCtx);
   if (auto Error = BinaryOrError.takeError()) {
     return Error;
   }
   KernelBinary *Binary = *BinaryOrError;
-  RTCBundleBinaryInfo BBI{BinaryFormat::SPIRV,
-                          Mod.getDataLayout().getPointerSizeInBits(),
-                          Binary->address(), Binary->size()};
-  return BBI;
+  RTCDevImgBinaryInfo DIBI{BinaryFormat::SPIRV,
+                           Mod.getDataLayout().getPointerSizeInBits(),
+                           Binary->address(), Binary->size()};
+  return DIBI;
 }
 
 llvm::Expected<KernelBinary *>

sycl-jit/jit-compiler/lib/translation/KernelTranslation.h

@@ -27,8 +27,8 @@ class KernelTranslator {
   static llvm::Error translateKernel(SYCLKernelInfo &Kernel, llvm::Module &Mod,
                                      JITContext &JITCtx, BinaryFormat Format);
 
-  static llvm::Expected<RTCBundleBinaryInfo>
-  translateBundleToSPIRV(llvm::Module &Mod, JITContext &JITCtx);
+  static llvm::Expected<RTCDevImgBinaryInfo>
+  translateDevImgToSPIRV(llvm::Module &Mod, JITContext &JITCtx);
 
 private:
   ///