Support multiple device images per RTC bundle.

Signed-off-by: Julian Oppermann <[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

@@ -232,8 +232,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());
@@ -435,15 +433,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
@@ -486,70 +504,83 @@ llvm::Expected<PostLinkResult> jit_compiler::performPostLink(
       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. due to 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>
@@ -606,21 +637,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:
   ///