From 7e92fb92ec21bb58d7ee4df6b4d2b295a598e53f Mon Sep 17 00:00:00 2001
From: Nikita Sirgienko <nikita.sirgienko@intel.com>
Date: Thu, 20 Apr 2023 20:12:30 +0200
Subject: [PATCH] Cycles: oneAPI: Fix kernels preloading in case of
 incompatible AoT binaries

When running oneAPI with AoT binaries, on hardware that's not compatible with
these, recompilation could have been missing from the kernels loading phase and
happen during execution instead.

These changes fixes it, any kernel compilation will now happen during the
kernels loading phase.
---
 intern/cycles/kernel/CMakeLists.txt           |  2 --
 intern/cycles/kernel/device/oneapi/kernel.cpp | 36 +++++++++++--------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/intern/cycles/kernel/CMakeLists.txt b/intern/cycles/kernel/CMakeLists.txt
index f3c970fd8bf..71466ace54c 100644
--- a/intern/cycles/kernel/CMakeLists.txt
+++ b/intern/cycles/kernel/CMakeLists.txt
@@ -778,8 +778,6 @@ if(WITH_CYCLES_DEVICE_ONEAPI)
 
   # Host execution won't use GPU binaries, no need to compile them.
   if(WITH_CYCLES_ONEAPI_BINARIES AND NOT WITH_CYCLES_ONEAPI_HOST_TASK_EXECUTION)
-    # AoT binaries aren't currently reused when calling sycl::build.
-    list(APPEND sycl_compiler_flags -DWITH_CYCLES_ONEAPI_BINARIES)
     # Iterate over all targest and their options
     list(JOIN CYCLES_ONEAPI_SYCL_TARGETS "," targets_string)
     list(APPEND sycl_compiler_flags -fsycl-targets=${targets_string})
diff --git a/intern/cycles/kernel/device/oneapi/kernel.cpp b/intern/cycles/kernel/device/oneapi/kernel.cpp
index 6d82b79c59e..e4e0f01add1 100644
--- a/intern/cycles/kernel/device/oneapi/kernel.cpp
+++ b/intern/cycles/kernel/device/oneapi/kernel.cpp
@@ -174,6 +174,14 @@ bool oneapi_kernel_is_required_for_features(const std::string &kernel_name,
   return true;
 }
 
+bool oneapi_kernel_is_raytrace_or_mnee(const std::string &kernel_name)
+{
+  return (kernel_name.find(device_kernel_as_string(DEVICE_KERNEL_INTEGRATOR_SHADE_SURFACE_MNEE)) !=
+          std::string::npos) ||
+         (kernel_name.find(device_kernel_as_string(
+              DEVICE_KERNEL_INTEGRATOR_SHADE_SURFACE_RAYTRACE)) != std::string::npos);
+}
+
 bool oneapi_kernel_is_using_embree(const std::string &kernel_name)
 {
 #  ifdef WITH_EMBREE_GPU
@@ -182,8 +190,7 @@ bool oneapi_kernel_is_using_embree(const std::string &kernel_name)
     DeviceKernel kernel = (DeviceKernel)i;
     if (device_kernel_has_intersection(kernel)) {
       if (kernel_name.find(device_kernel_as_string(kernel)) != std::string::npos) {
-        return !(kernel == DEVICE_KERNEL_INTEGRATOR_SHADE_SURFACE_RAYTRACE ||
-                 kernel == DEVICE_KERNEL_INTEGRATOR_SHADE_SURFACE_MNEE);
+        return !oneapi_kernel_is_raytrace_or_mnee(kernel_name);
       }
     }
   }
@@ -241,10 +248,6 @@ bool oneapi_load_kernels(SyclQueue *queue_,
   }
 #  endif
 
-#  ifdef WITH_CYCLES_ONEAPI_BINARIES
-  (void)queue_;
-  (void)kernel_features;
-#  else
   try {
     sycl::kernel_bundle<sycl::bundle_state::input> all_kernels_bundle =
         sycl::get_kernel_bundle<sycl::bundle_state::input>(queue->get_context(),
@@ -260,18 +263,22 @@ bool oneapi_load_kernels(SyclQueue *queue_,
         continue;
       }
 
-      sycl::kernel_bundle<sycl::bundle_state::input> one_kernel_bundle_input =
-          sycl::get_kernel_bundle<sycl::bundle_state::input>(queue->get_context(), {kernel_id});
-#    ifdef WITH_EMBREE_GPU
-      /* This is expected to be the default, we set it again to be sure. */
-      if (one_kernel_bundle_input
-              .has_specialization_constant<ONEAPIKernelContext::oneapi_embree_features>()) {
+#  ifdef WITH_EMBREE_GPU
+      if (oneapi_kernel_is_using_embree(kernel_name) ||
+          oneapi_kernel_is_raytrace_or_mnee(kernel_name)) {
+        sycl::kernel_bundle<sycl::bundle_state::input> one_kernel_bundle_input =
+            sycl::get_kernel_bundle<sycl::bundle_state::input>(queue->get_context(), {kernel_id});
         one_kernel_bundle_input
             .set_specialization_constant<ONEAPIKernelContext::oneapi_embree_features>(
                 RTC_FEATURE_FLAG_NONE);
+        sycl::build(one_kernel_bundle_input);
+        continue;
       }
-#    endif
-      sycl::build(one_kernel_bundle_input);
+#  endif
+      /* This call will ensure that AoT or cached JIT binaries are available
+       * for execution. It will trigger compilation if it is not already the case. */
+      (void)sycl::get_kernel_bundle<sycl::bundle_state::executable>(queue->get_context(),
+                                                                    {kernel_id});
     }
   }
   catch (sycl::exception const &e) {
@@ -280,7 +287,6 @@ bool oneapi_load_kernels(SyclQueue *queue_,
     }
     return false;
   }
-#  endif
   return true;
 }