Cycles: Add support for parallel compilation of OptiX module

OptiX 7.4 adds support for splitting the costly creation of an OptiX
module into smaller tasks that can be executed in parallel on a
thread pool.
This is only really relevant for the "shader_raytrace" kernel variant
as the main one is small and compiles fast either way. It sheds of
a few seconds there (total gain is not massive currently, since it is
difficult for the compiler to split up the huge shading entry point
that is the primary one taking up time, but it is still measurable).

Differential Revision: https://developer.blender.org/D14845

1 files changed, 37 insertions, 0 deletions

diff --git a/intern/cycles/device/optix/device_impl.cpp b/intern/cycles/device/optix/device_impl.cpp
index 8830d8c44ac..6329144131e 100644
--- a/intern/cycles/device/optix/device_impl.cpp
+++ b/intern/cycles/device/optix/device_impl.cpp
@@ -23,6 +23,7 @@
 #  include "util/md5.h"
 #  include "util/path.h"
 #  include "util/progress.h"
+#  include "util/task.h"
 #  include "util/time.h"
 
 #  undef __KERNEL_CPU__
@@ -216,6 +217,24 @@ static OptixResult optixUtilDenoiserInvokeTiled(OptixDenoiser denoiser,
   return OPTIX_SUCCESS;
 }
 
+#  if OPTIX_ABI_VERSION >= 55
+static void execute_optix_task(TaskPool &pool, OptixTask task, OptixResult &failure_reason)
+{
+  OptixTask additional_tasks[16];
+  unsigned int num_additional_tasks = 0;
+
+  const OptixResult result = optixTaskExecute(task, additional_tasks, 16, &num_additional_tasks);
+  if (result == OPTIX_SUCCESS) {
+    for (unsigned int i = 0; i < num_additional_tasks; ++i) {
+      pool.push(function_bind(&execute_optix_task, std::ref(pool), additional_tasks[i], std::ref(failure_reason)));
+    }
+  }
+  else {
+    failure_reason = result;
+  }
+}
+#  endif
+
 }  // namespace
 
 OptiXDevice::Denoiser::Denoiser(OptiXDevice *device)
@@ -453,6 +472,23 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
       return false;
     }
 
+#  if OPTIX_ABI_VERSION >= 55
+    OptixTask task = nullptr;
+    OptixResult result = optixModuleCreateFromPTXWithTasks(context,
+                                                           &module_options,
+                                                           &pipeline_options,
+                                                           ptx_data.data(),
+                                                           ptx_data.size(),
+                                                           nullptr,
+                                                           nullptr,
+                                                           &optix_module,
+                                                           &task);
+    if (result == OPTIX_SUCCESS) {
+      TaskPool pool;
+      execute_optix_task(pool, task, result);
+      pool.wait_work();
+    }
+#  else
     const OptixResult result = optixModuleCreateFromPTX(context,
                                                         &module_options,
                                                         &pipeline_options,
@@ -461,6 +497,7 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
                                                         nullptr,
                                                         0,
                                                         &optix_module);
+#  endif
     if (result != OPTIX_SUCCESS) {
       set_error(string_printf("Failed to load OptiX kernel from '%s' (%s)",
                               ptx_filename.c_str(),