[SYCL][RTC] Extend device code split E2E test (#16685)

Extends test to cover all explicit device code split modes, and adds a
test for implicit device code splitting.

---------

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

Author: jopperm

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp

@@ -82,6 +82,18 @@ void vector_add_esimd(float *A, float *B, float *C) {
     }
 )===";
 
+auto constexpr DeviceCodeSplitSource = R"===(
+#include <sycl/sycl.hpp>
+
+template<typename T, unsigned WG = 16> SYCL_EXTERNAL 
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY(sycl::ext::oneapi::experimental::nd_range_kernel<1>)
+[[sycl::reqd_work_group_size(WG)]]
+void vec_add(T* in1, T* in2, T* out){
+  size_t id = sycl::ext::oneapi::this_work_item::get_nd_item<1>().get_global_linear_id();
+  out[id] = in1[id] + in2[id];
+}
+)===";
+
 auto constexpr BadSource = R"===(
 #include <sycl/sycl.hpp>
 
@@ -206,12 +218,7 @@ int test_build_and_run() {
       ctx, syclex::source_language::sycl_jit, SYCLSource,
       syclex::properties{incFiles2});
 
-  exe_kb kbExe3 = syclex::build(
-      kbSrc2, syclex::properties{
-                  syclex::build_options{"-fsycl-device-code-split=per_kernel"},
-                  syclex::registered_kernel_names{"ff_templated<int>"}});
-  assert(std::distance(kbExe3.begin(), kbExe3.end()) == 2 &&
-         "Expected 2 device images");
+  exe_kb kbExe3 = syclex::build(kbSrc2);
   sycl::kernel k3 = kbExe3.ext_oneapi_get_kernel("ff_cp");
   test_1(q, k3, 37 + 7);
 
@@ -222,6 +229,58 @@ int test_build_and_run() {
   return 0;
 }
 
+int test_device_code_split() {
+  namespace syclex = sycl::ext::oneapi::experimental;
+  using source_kb = sycl::kernel_bundle<sycl::bundle_state::ext_oneapi_source>;
+  using exe_kb = sycl::kernel_bundle<sycl::bundle_state::executable>;
+
+  sycl::queue q;
+  sycl::context ctx = q.get_context();
+
+  bool ok =
+      q.get_device().ext_oneapi_can_compile(syclex::source_language::sycl_jit);
+  if (!ok) {
+    std::cout << "Apparently this device does not support `sycl_jit` source "
+                 "kernel bundle extension: "
+              << q.get_device().get_info<sycl::info::device::name>()
+              << std::endl;
+    return -1;
+  }
+
+  source_kb kbSrc = syclex::create_kernel_bundle_from_source(
+      ctx, syclex::source_language::sycl_jit, DeviceCodeSplitSource);
+
+  // Test explicit device code split
+  std::vector<std::string> names{"vec_add<float>", "vec_add<int>",
+                                 "vec_add<short>"};
+  auto build = [&](const std::string &mode) -> size_t {
+    exe_kb kbExe = syclex::build(
+        kbSrc, syclex::properties{
+                   syclex::registered_kernel_names{names},
+                   syclex::build_options{"-fsycl-device-code-split=" + mode}});
+    return std::distance(kbExe.begin(), kbExe.end());
+  };
+
+  size_t perKernelNImg = build("per_kernel");
+  size_t perSourceNImg = build("per_source");
+  size_t offNImg = build("off");
+  size_t autoNImg = build("auto");
+
+  assert(perKernelNImg == 3);
+  assert(perSourceNImg == 1);
+  assert(offNImg == 1);
+  assert(autoNImg >= offNImg && autoNImg <= perKernelNImg);
+
+  // Test implicit device code split
+  names = {"vec_add<float, 8>", "vec_add<float, 16>"};
+  exe_kb kbDiffWorkGroupSizes = syclex::build(
+      kbSrc, syclex::properties{syclex::registered_kernel_names{names}});
+  assert(std::distance(kbDiffWorkGroupSizes.begin(),
+                       kbDiffWorkGroupSizes.end()) == 2);
+
+  return 0;
+}
+
 int test_esimd() {
   namespace syclex = sycl::ext::oneapi::experimental;
   using source_kb = sycl::kernel_bundle<sycl::bundle_state::ext_oneapi_source>;
@@ -393,8 +452,8 @@ int test_warning() {
 int main(int argc, char **) {
 #ifdef SYCL_EXT_ONEAPI_KERNEL_COMPILER
   int optional_tests = (argc > 1) ? test_warning() : 0;
-  return test_build_and_run() || test_esimd() || test_unsupported_options() ||
-         test_error() || optional_tests;
+  return test_build_and_run() || test_device_code_split() || test_esimd() ||
+         test_unsupported_options() || test_error() || optional_tests;
 #else
   static_assert(false, "Kernel Compiler feature test macro undefined");
 #endif