[SYCL][RTC] Extend mixed SYCL/ESIMD support (intel#17373)

Add support for the situation that after the first round of splitting, a
`ModuleDesc` contains a mix of SYCL and ESIMD kernels.

---------

Signed-off-by: Julian Oppermann <[email protected]>

Author: jopperm

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

@@ -659,79 +659,82 @@ jit_compiler::performPostLink(std::unique_ptr<llvm::Module> Module,
   // TODO: Call `verifyNoCrossModuleDeviceGlobalUsage` if device globals shall
   //       be processed.
 
-  // TODO: This allocation assumes that there are no further splits required,
-  //       i.e. there are no mixed SYCL/ESIMD modules.
-  RTCBundleInfo BundleInfo;
-  BundleInfo.DevImgInfos = DynArray<RTCDevImgInfo>{Splitter->remainingSplits()};
+  SmallVector<RTCDevImgInfo> DevImgInfoVec;
   SmallVector<std::unique_ptr<llvm::Module>> Modules;
 
-  auto *DevImgInfoIt = BundleInfo.DevImgInfos.begin();
-  while (Splitter->hasMoreSplits()) {
-    assert(DevImgInfoIt != BundleInfo.DevImgInfos.end());
+  // TODO: The following logic is missing the ability to link ESIMD and SYCL
+  //       modules back together, which would be requested via
+  //       `-fno-sycl-device-code-split-esimd` as a prerequisite for compiling
+  //       `invoke_simd` code.
 
+  while (Splitter->hasMoreSplits()) {
     ModuleDesc MDesc = Splitter->nextSplit();
-    RTCDevImgInfo &DevImgInfo = *DevImgInfoIt++;
 
     // TODO: Call `MDesc.fixupLinkageOfDirectInvokeSimdTargets()` when
     //       `invoke_simd` is supported.
 
     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()};
+    for (auto &ES : ESIMDSplits) {
+      MDesc = std::move(ES);
+
+      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);
       }
-    };
 
-    Modules.push_back(MDesc.releaseModulePtr());
+      MDesc.saveSplitInformationAsMetadata();
+
+      RTCDevImgInfo &DevImgInfo = DevImgInfoVec.emplace_back();
+      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());
+    }
   }
 
+  assert(DevImgInfoVec.size() == Modules.size());
+  RTCBundleInfo BundleInfo;
+  BundleInfo.DevImgInfos = DynArray<RTCDevImgInfo>{DevImgInfoVec.size()};
+  std::move(DevImgInfoVec.begin(), DevImgInfoVec.end(),
+            BundleInfo.DevImgInfos.begin());
+
   return PostLinkResult{std::move(BundleInfo), std::move(Modules)};
 }
 

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp

@@ -515,6 +515,18 @@ int test_esimd() {
   run_2(q, kESIMD, true, 2.38f);
   run_2(q, kSYCL, false, 1.41f);
 
+  // Deactivate implicit module splitting to exercise the downstream
+  // ESIMD-specific splitting.
+  source_kb kbSrcMixed2 = syclex::create_kernel_bundle_from_source(
+      ctx, syclex::source_language::sycl_jit, mixedSource);
+  exe_kb kbExeMixed2 =
+      syclex::build(kbSrcMixed2, syclex::properties{syclex::build_options{
+                                     "-fsycl-device-code-split=off"}});
+
+  assert(kbExeMixed2.ext_oneapi_has_kernel("vector_add_esimd"));
+  assert(kbExeMixed2.ext_oneapi_has_kernel("vec_add"));
+  assert(std::distance(kbExeMixed2.begin(), kbExeMixed2.end()) == 2);
+
   return 0;
 }