Cycles: Split OptiX kernels into smaller modules to improve load time

This helps especially for OSL, where load time is very long. By
splitting off shader raytrace, MNEE and volumes, the first time
rendering is much faster.

The downside is that noinline functions will be duplicated. However OSL
startup performance is very bad currently and this seems the better
trade-off for now. There are ways to make this work if we do not mark
noinline functions as static, but this will require some bigger code
reorganization.

Without OSL, shader raytrace and MNEE were combined in a single module,
and that has been split into two as well.

Besides a better user experience, This will fix OptiX OSL test timeouts
on the buildbot, where some tests need to compute a few different
specializations and the 600s timeout is exceeded, with the longest test
run time being around 300s.

Pull Request: https://projects.blender.org/blender/blender/pulls/159499

intern/cycles/device/optix/device_impl.cpp

@@ -113,6 +113,12 @@ OptiXDevice::~OptiXDevice()
   if (optix_module != nullptr) {
     optixModuleDestroy(optix_module);
   }
+  if (mnee_module != nullptr) {
+    optixModuleDestroy(mnee_module);
+  }
+  if (shader_raytrace_module != nullptr) {
+    optixModuleDestroy(shader_raytrace_module);
+  }
   for (int i = 0; i < 2; ++i) {
     if (builtin_modules[i] != nullptr) {
       optixModuleDestroy(builtin_modules[i]);
@@ -133,6 +139,9 @@ OptiXDevice::~OptiXDevice()
   if (osl_camera_module != nullptr) {
     optixModuleDestroy(osl_camera_module);
   }
+  if (osl_volume_module != nullptr) {
+    optixModuleDestroy(osl_volume_module);
+  }
   for (const OptixModule &module : osl_modules) {
     if (module != nullptr) {
       optixModuleDestroy(module);
@@ -221,13 +230,16 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   }
 
 #  ifdef WITH_OSL
-  /* TODO: Consider splitting kernels into an OSL-camera-only and a full-OSL variant. */
   const bool use_osl_shading = (kernel_features & KERNEL_FEATURE_OSL_SHADING);
   const bool use_osl_camera = (kernel_features & KERNEL_FEATURE_OSL_CAMERA);
+  const bool use_osl_volume = use_osl_shading && (kernel_features & KERNEL_FEATURE_VOLUME);
 #  else
   const bool use_osl_shading = false;
   const bool use_osl_camera = false;
+  const bool use_osl_volume = false;
 #  endif
+  const bool use_shader_raytrace = (kernel_features & KERNEL_FEATURE_NODE_RAYTRACE);
+  const bool use_mnee = (kernel_features & KERNEL_FEATURE_MNEE);
 
   /* Skip creating OptiX module if only doing denoising. */
   const bool need_optix_kernels = (kernel_features &
@@ -236,10 +248,7 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   /* Detect existence of OptiX kernel and SDK here early. So we can error out
    * before compiling the CUDA kernels, to avoid failing right after when
    * compiling the OptiX kernel. */
-  string suffix = use_osl_shading ? "_osl" :
-                  (kernel_features & (KERNEL_FEATURE_NODE_RAYTRACE | KERNEL_FEATURE_MNEE)) ?
-                                    "_shader_raytrace" :
-                                    "";
+  string suffix = use_osl_shading ? "_osl" : "";
   string ptx_filename;
   if (need_optix_kernels) {
     ptx_filename = path_get("lib/kernel_optix" + suffix + ".ptx.zst");
@@ -278,6 +287,14 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
     optixModuleDestroy(optix_module);
     optix_module = nullptr;
   }
+  if (mnee_module != nullptr) {
+    optixModuleDestroy(mnee_module);
+    mnee_module = nullptr;
+  }
+  if (shader_raytrace_module != nullptr) {
+    optixModuleDestroy(shader_raytrace_module);
+    shader_raytrace_module = nullptr;
+  }
   for (int i = 0; i < 2; ++i) {
     if (builtin_modules[i] != nullptr) {
       optixModuleDestroy(builtin_modules[i]);
@@ -302,6 +319,10 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
     optixModuleDestroy(osl_camera_module);
     osl_camera_module = nullptr;
   }
+  if (osl_volume_module != nullptr) {
+    optixModuleDestroy(osl_volume_module);
+    osl_volume_module = nullptr;
+  }
 
   /* Recreating base OptiX module invalidates all OSL modules too, since they link against it. */
   for (const OptixModule &module : osl_modules) {
@@ -364,27 +385,113 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
     pipeline_options.traversableGraphFlags = OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_ANY;
   }
 
-  { /* Load and compile PTX module with OptiX kernels. */
-    string ptx_data;
+  { /* Load and compile PTX modules with OptiX kernels.
+     * All in one TaskPool so the OptiX driver can compile them concurrently. */
+    string base_ptx_data;
     if (use_adaptive_compilation() || path_file_size(ptx_filename) == -1) {
       string cflags = compile_kernel_get_common_cflags(kernel_features);
       ptx_filename = compile_kernel(cflags, ("kernel" + suffix).c_str(), true);
     }
-    if (ptx_filename.empty() || !path_read_compressed_text(ptx_filename, ptx_data)) {
+    if (ptx_filename.empty() || !path_read_compressed_text(ptx_filename, base_ptx_data)) {
       set_error(string_printf("Failed to load OptiX kernel from '%s'", ptx_filename.c_str()));
       return false;
     }
 
-    TaskPool pool;
-    OptixResult result;
-    create_optix_module(pool, module_options, ptx_data, optix_module, result);
-    pool.wait_work();
-    if (result != OPTIX_SUCCESS) {
-      set_error(string_printf("Failed to load OptiX kernel from '%s' (%s)",
-                              ptx_filename.c_str(),
-                              optixGetErrorName(result)));
+    auto load_optional_module = [this](const string &name, string &ptx_data) -> bool {
+      const string filename = path_get("lib/" + name + ".ptx.zst");
+      if (!path_read_compressed_text(filename, ptx_data)) {
+        set_error(string_printf("Failed to load OptiX kernel from '%s'", filename.c_str()));
+        return false;
+      }
+      return true;
+    };
+
+    string mnee_ptx_data;
+    if (use_mnee &&
+        !load_optional_module(use_osl_shading ? "kernel_optix_osl_mnee" : "kernel_optix_mnee",
+                              mnee_ptx_data))
+    {
       return false;
     }
+    string shader_raytrace_ptx_data;
+    if (use_shader_raytrace &&
+        !load_optional_module(use_osl_shading ? "kernel_optix_osl_shader_raytrace" :
+                                                "kernel_optix_shader_raytrace",
+                              shader_raytrace_ptx_data))
+    {
+      return false;
+    }
+#  ifdef WITH_OSL
+    string osl_camera_ptx_data;
+    if (use_osl_camera && !load_optional_module("kernel_optix_osl_camera", osl_camera_ptx_data)) {
+      return false;
+    }
+    string osl_volume_ptx_data;
+    if (use_osl_volume && !load_optional_module("kernel_optix_osl_volume", osl_volume_ptx_data)) {
+      return false;
+    }
+#  endif
+
+    TaskPool pool;
+    OptixResult base_result = OPTIX_SUCCESS;
+    OptixResult mnee_result = OPTIX_SUCCESS;
+    OptixResult shader_raytrace_result = OPTIX_SUCCESS;
+#  ifdef WITH_OSL
+    OptixResult osl_camera_result = OPTIX_SUCCESS;
+    OptixResult osl_volume_result = OPTIX_SUCCESS;
+#  endif
+
+    create_optix_module(pool, module_options, base_ptx_data, optix_module, base_result);
+    if (use_mnee) {
+      create_optix_module(pool, module_options, mnee_ptx_data, mnee_module, mnee_result);
+    }
+    if (use_shader_raytrace) {
+      create_optix_module(pool,
+                          module_options,
+                          shader_raytrace_ptx_data,
+                          shader_raytrace_module,
+                          shader_raytrace_result);
+    }
+#  ifdef WITH_OSL
+    if (use_osl_camera) {
+      create_optix_module(
+          pool, module_options, osl_camera_ptx_data, osl_camera_module, osl_camera_result);
+    }
+    if (use_osl_volume) {
+      create_optix_module(
+          pool, module_options, osl_volume_ptx_data, osl_volume_module, osl_volume_result);
+    }
+#  endif
+    pool.wait_work();
+
+    if (base_result != OPTIX_SUCCESS) {
+      set_error(string_printf("Failed to load OptiX kernel from '%s' (%s)",
+                              ptx_filename.c_str(),
+                              optixGetErrorName(base_result)));
+      return false;
+    }
+    if (mnee_result != OPTIX_SUCCESS) {
+      set_error(
+          string_printf("Failed to load OptiX MNEE kernel (%s)", optixGetErrorName(mnee_result)));
+      return false;
+    }
+    if (shader_raytrace_result != OPTIX_SUCCESS) {
+      set_error(string_printf("Failed to load OptiX shader ray-tracing kernel (%s)",
+                              optixGetErrorName(shader_raytrace_result)));
+      return false;
+    }
+#  ifdef WITH_OSL
+    if (osl_camera_result != OPTIX_SUCCESS) {
+      set_error(string_printf("Failed to load OptiX OSL camera kernel (%s)",
+                              optixGetErrorName(osl_camera_result)));
+      return false;
+    }
+    if (osl_volume_result != OPTIX_SUCCESS) {
+      set_error(string_printf("Failed to load OptiX OSL volume kernel (%s)",
+                              optixGetErrorName(osl_volume_result)));
+      return false;
+    }
+#  endif
   }
 
   /* Create program groups. */
@@ -527,9 +634,9 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   }
 
   /* Shader ray-tracing replaces some functions with direct callables. */
-  if (kernel_features & KERNEL_FEATURE_NODE_RAYTRACE) {
+  if (use_shader_raytrace) {
     group_descs[PG_RGEN_SHADE_SURFACE_RAYTRACE].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
-    group_descs[PG_RGEN_SHADE_SURFACE_RAYTRACE].raygen.module = optix_module;
+    group_descs[PG_RGEN_SHADE_SURFACE_RAYTRACE].raygen.module = shader_raytrace_module;
     group_descs[PG_RGEN_SHADE_SURFACE_RAYTRACE].raygen.entryFunctionName =
         "__raygen__kernel_optix_integrator_shade_surface_raytrace";
 
@@ -537,18 +644,18 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
      * there. */
     if (!use_osl_shading) {
       group_descs[PG_CALL_SVM_AO].kind = OPTIX_PROGRAM_GROUP_KIND_CALLABLES;
-      group_descs[PG_CALL_SVM_AO].callables.moduleDC = optix_module;
+      group_descs[PG_CALL_SVM_AO].callables.moduleDC = shader_raytrace_module;
       group_descs[PG_CALL_SVM_AO].callables.entryFunctionNameDC = "__direct_callable__svm_node_ao";
       group_descs[PG_CALL_SVM_BEVEL].kind = OPTIX_PROGRAM_GROUP_KIND_CALLABLES;
-      group_descs[PG_CALL_SVM_BEVEL].callables.moduleDC = optix_module;
+      group_descs[PG_CALL_SVM_BEVEL].callables.moduleDC = shader_raytrace_module;
       group_descs[PG_CALL_SVM_BEVEL].callables.entryFunctionNameDC =
           "__direct_callable__svm_node_bevel";
     }
   }
 
-  if (kernel_features & KERNEL_FEATURE_MNEE) {
+  if (use_mnee) {
     group_descs[PG_RGEN_INTERSECT_MNEE].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
-    group_descs[PG_RGEN_INTERSECT_MNEE].raygen.module = optix_module;
+    group_descs[PG_RGEN_INTERSECT_MNEE].raygen.module = mnee_module;
     group_descs[PG_RGEN_INTERSECT_MNEE].raygen.entryFunctionName =
         "__raygen__kernel_optix_integrator_intersect_mnee";
   }
@@ -571,14 +678,16 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
     group_descs[PG_RGEN_SHADE_SURFACE].raygen.module = optix_module;
     group_descs[PG_RGEN_SHADE_SURFACE].raygen.entryFunctionName =
         "__raygen__kernel_optix_integrator_shade_surface";
-    group_descs[PG_RGEN_SHADE_VOLUME].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
-    group_descs[PG_RGEN_SHADE_VOLUME].raygen.module = optix_module;
-    group_descs[PG_RGEN_SHADE_VOLUME].raygen.entryFunctionName =
-        "__raygen__kernel_optix_integrator_shade_volume";
-    group_descs[PG_RGEN_SHADE_VOLUME_RAY_MARCHING].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
-    group_descs[PG_RGEN_SHADE_VOLUME_RAY_MARCHING].raygen.module = optix_module;
-    group_descs[PG_RGEN_SHADE_VOLUME_RAY_MARCHING].raygen.entryFunctionName =
-        "__raygen__kernel_optix_integrator_shade_volume_ray_marching";
+    if (use_osl_volume) {
+      group_descs[PG_RGEN_SHADE_VOLUME].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
+      group_descs[PG_RGEN_SHADE_VOLUME].raygen.module = osl_volume_module;
+      group_descs[PG_RGEN_SHADE_VOLUME].raygen.entryFunctionName =
+          "__raygen__kernel_optix_integrator_shade_volume";
+      group_descs[PG_RGEN_SHADE_VOLUME_RAY_MARCHING].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
+      group_descs[PG_RGEN_SHADE_VOLUME_RAY_MARCHING].raygen.module = osl_volume_module;
+      group_descs[PG_RGEN_SHADE_VOLUME_RAY_MARCHING].raygen.entryFunctionName =
+          "__raygen__kernel_optix_integrator_shade_volume_ray_marching";
+    }
     group_descs[PG_RGEN_SHADE_SHADOW].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
     group_descs[PG_RGEN_SHADE_SHADOW].raygen.module = optix_module;
     group_descs[PG_RGEN_SHADE_SHADOW].raygen.entryFunctionName =
@@ -608,25 +717,6 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
 #  ifdef WITH_OSL
   /* When using custom OSL cameras, integrator_init_from_camera is its own specialized module. */
   if (use_osl_camera) {
-    /* Load and compile the OSL camera PTX module. */
-    string ptx_data, ptx_filename = path_get("lib/kernel_optix_osl_camera.ptx.zst");
-    if (!path_read_compressed_text(ptx_filename, ptx_data)) {
-      set_error(
-          string_printf("Failed to load OptiX OSL camera kernel from '%s'", ptx_filename.c_str()));
-      return false;
-    }
-
-    TaskPool pool;
-    OptixResult result;
-    create_optix_module(pool, module_options, ptx_data, osl_camera_module, result);
-    pool.wait_work();
-    if (result != OPTIX_SUCCESS) {
-      set_error(string_printf("Failed to load OptiX kernel from '%s' (%s)",
-                              ptx_filename.c_str(),
-                              optixGetErrorName(result)));
-      return false;
-    }
-
     group_descs[PG_RGEN_INIT_FROM_CAMERA].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
     group_descs[PG_RGEN_INIT_FROM_CAMERA].raygen.module = osl_camera_module;
     group_descs[PG_RGEN_INIT_FROM_CAMERA].raygen.entryFunctionName =
@@ -638,10 +728,18 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
       context, group_descs, NUM_PROGRAM_GROUPS, &group_options, nullptr, nullptr, groups));
 
   /* Get program stack sizes. */
-  auto get_pipeline_stack_size = [&](OptixPipeline pipeline, unsigned int &trace_css) {
+  auto get_pipeline_stack_size = [&](OptixPipeline pipeline,
+                                     const vector<OptixProgramGroup> &pipeline_groups,
+                                     unsigned int &trace_css) {
     vector<OptixStackSizes> stack_size(NUM_PROGRAM_GROUPS);
     for (int i = 0; i < NUM_PROGRAM_GROUPS; ++i) {
-      optix_assert(optixProgramGroupGetStackSize(groups[i], &stack_size[i], pipeline));
+      /* Only groups that are part of the pipeline, otherwise this is an error. */
+      if (groups[i] != nullptr &&
+          std::find(pipeline_groups.begin(), pipeline_groups.end(), groups[i]) !=
+              pipeline_groups.end())
+      {
+        optix_assert(optixProgramGroupGetStackSize(groups[i], &stack_size[i], pipeline));
+      }
     }
 
     /* Calculate maximum trace continuation stack size. */
@@ -685,7 +783,9 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   sbt_data.alloc(NUM_PROGRAM_GROUPS);
   memset(sbt_data.host_pointer, 0, sizeof(SbtRecord) * NUM_PROGRAM_GROUPS);
   for (int i = 0; i < NUM_PROGRAM_GROUPS; ++i) {
-    optix_assert(optixSbtRecordPackHeader(groups[i], &sbt_data[i]));
+    if (groups[i] != nullptr) {
+      optix_assert(optixSbtRecordPackHeader(groups[i], &sbt_data[i]));
+    }
   }
   sbt_data.copy_to_device(); /* Upload SBT to device. */
 
@@ -753,7 +853,8 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
                                      &pipelines[PIP_SHADE]));
 
     unsigned int trace_css;
-    vector<OptixStackSizes> stack_size = get_pipeline_stack_size(pipelines[PIP_SHADE], trace_css);
+    vector<OptixStackSizes> stack_size = get_pipeline_stack_size(
+        pipelines[PIP_SHADE], pipeline_groups, trace_css);
 
     /* Combine ray generation and trace continuation stack size. */
     const unsigned int css = std::max(stack_size[PG_RGEN_SHADE_SURFACE_RAYTRACE].cssRG,
@@ -811,8 +912,8 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
                                      &pipelines[PIP_INTERSECT]));
 
     unsigned int trace_css;
-    vector<OptixStackSizes> stack_size = get_pipeline_stack_size(pipelines[PIP_INTERSECT],
-                                                                 trace_css);
+    vector<OptixStackSizes> stack_size = get_pipeline_stack_size(
+        pipelines[PIP_INTERSECT], pipeline_groups, trace_css);
 
     /* Calculate continuation stack size based on the maximum of all ray generation stack sizes. */
     const unsigned int css =
@@ -1035,7 +1136,9 @@ bool OptiXDevice::load_osl_kernels()
   /* Update SBT with new entries. */
   sbt_data.alloc(NUM_PROGRAM_GROUPS + osl_groups.size());
   for (int i = 0; i < NUM_PROGRAM_GROUPS; ++i) {
-    optix_assert(optixSbtRecordPackHeader(groups[i], &sbt_data[i]));
+    if (groups[i] != nullptr) {
+      optix_assert(optixSbtRecordPackHeader(groups[i], &sbt_data[i]));
+    }
   }
   for (size_t i = 0; i < osl_groups.size(); ++i) {
     if (osl_groups[i] != nullptr) {
@@ -1060,11 +1163,21 @@ bool OptiXDevice::load_osl_kernels()
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_LIGHT_NEE]);
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_LIGHT_FORWARD]);
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_SURFACE]);
-    pipeline_groups.push_back(groups[PG_RGEN_SHADE_SURFACE_RAYTRACE]);
-    pipeline_groups.push_back(groups[PG_CALL_SVM_AO]);
-    pipeline_groups.push_back(groups[PG_CALL_SVM_BEVEL]);
-    pipeline_groups.push_back(groups[PG_RGEN_INTERSECT_MNEE]);
-    pipeline_groups.push_back(groups[PG_RGEN_SHADE_VOLUME]);
+    if (groups[PG_RGEN_SHADE_SURFACE_RAYTRACE] != nullptr) {
+      pipeline_groups.push_back(groups[PG_RGEN_SHADE_SURFACE_RAYTRACE]);
+    }
+    if (groups[PG_CALL_SVM_AO] != nullptr) {
+      pipeline_groups.push_back(groups[PG_CALL_SVM_AO]);
+    }
+    if (groups[PG_CALL_SVM_BEVEL] != nullptr) {
+      pipeline_groups.push_back(groups[PG_CALL_SVM_BEVEL]);
+    }
+    if (groups[PG_RGEN_INTERSECT_MNEE] != nullptr) {
+      pipeline_groups.push_back(groups[PG_RGEN_INTERSECT_MNEE]);
+    }
+    if (groups[PG_RGEN_SHADE_VOLUME] != nullptr) {
+      pipeline_groups.push_back(groups[PG_RGEN_SHADE_VOLUME]);
+    }
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_SHADOW]);
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_DEDICATED_LIGHT]);
     pipeline_groups.push_back(groups[PG_RGEN_EVAL_DISPLACE]);
@@ -1093,7 +1206,10 @@ bool OptiXDevice::load_osl_kernels()
     vector<OptixStackSizes> osl_stack_size(osl_groups.size());
 
     for (int i = 0; i < NUM_PROGRAM_GROUPS; ++i) {
-      optix_assert(optixProgramGroupGetStackSize(groups[i], &stack_size[i], pipelines[PIP_SHADE]));
+      if (groups[i] != nullptr) {
+        optix_assert(
+            optixProgramGroupGetStackSize(groups[i], &stack_size[i], pipelines[PIP_SHADE]));
+      }
     }
     for (size_t i = 0; i < osl_groups.size(); ++i) {
       if (osl_groups[i] != nullptr) {

intern/cycles/device/optix/device_impl.h

@@ -97,7 +97,9 @@ class OptiXDevice : public CUDADevice {
  public:
   OptixDeviceContext context = nullptr;
 
-  OptixModule optix_module = nullptr; /* All necessary OptiX kernels are in one module. */
+  OptixModule optix_module = nullptr;
+  OptixModule mnee_module = nullptr;
+  OptixModule shader_raytrace_module = nullptr;
   OptixModule builtin_modules[4] = {};
   OptixPipeline pipelines[NUM_PIPELINES] = {};
   OptixProgramGroup groups[NUM_PROGRAM_GROUPS] = {};
@@ -108,6 +110,7 @@ class OptiXDevice : public CUDADevice {
   vector<OptixModule> osl_modules;
   vector<OptixProgramGroup> osl_groups;
   OptixModule osl_camera_module = nullptr;
+  OptixModule osl_volume_module = nullptr;
   device_vector<uint8_t> osl_colorsystem;
 #  endif
 

intern/cycles/kernel/device/optix/CMakeLists.txt

@@ -12,6 +12,7 @@ set(INC_SYS
 
 set(SRC_KERNEL_DEVICE_OPTIX
   kernel.cu
+  kernel_mnee.cu
   kernel_shader_raytrace.cu
 )
 
@@ -23,6 +24,9 @@ if(WITH_CYCLES_OSL)
     ../../osl/services_optix.cu
     kernel_osl.cu
     kernel_osl_camera.cu
+    kernel_osl_mnee.cu
+    kernel_osl_shader_raytrace.cu
+    kernel_osl_volume.cu
   )
 endif()
 
@@ -105,6 +109,10 @@ if(WITH_CYCLES_CUDA_BINARIES AND WITH_CYCLES_DEVICE_OPTIX)
     kernel_optix
     "kernel.cu"
     "")
+  cycles_optix_kernel_add(
+    kernel_optix_mnee
+    "kernel_mnee.cu"
+    "")
   cycles_optix_kernel_add(
     kernel_optix_shader_raytrace
     "kernel_shader_raytrace.cu"
@@ -114,6 +122,18 @@ if(WITH_CYCLES_CUDA_BINARIES AND WITH_CYCLES_DEVICE_OPTIX)
       kernel_optix_osl
       "kernel_osl.cu"
       "--relocatable-device-code=true")
+    cycles_optix_kernel_add(
+      kernel_optix_osl_shader_raytrace
+      "kernel_osl_shader_raytrace.cu"
+      "--relocatable-device-code=true")
+    cycles_optix_kernel_add(
+      kernel_optix_osl_mnee
+      "kernel_osl_mnee.cu"
+      "--relocatable-device-code=true")
+    cycles_optix_kernel_add(
+      kernel_optix_osl_volume
+      "kernel_osl_volume.cu"
+      "--relocatable-device-code=true")
     cycles_optix_kernel_add(
       kernel_optix_osl_camera
       "kernel_osl_camera.cu"

intern/cycles/kernel/device/optix/kernel_mnee.cu

@@ -0,0 +1,19 @@
+/* SPDX-FileCopyrightText: 2011-2026 Blender Foundation
+ *
+ * SPDX-License-Identifier: Apache-2.0 */
+
+#include "kernel/device/optix/compat.h"
+#include "kernel/device/optix/globals.h"
+
+#include "kernel/device/gpu/image.h" /* Texture lookup uses normal CUDA intrinsics. */
+
+#include "kernel/integrator/intersect_mnee.h"
+
+extern "C" __global__ void __raygen__kernel_optix_integrator_intersect_mnee()
+{
+  const int global_index = optixGetLaunchIndex().x;
+  const int path_index = (kernel_params.path_index_array) ?
+                             kernel_params.path_index_array[global_index] :
+                             global_index;
+  integrator_intersect_mnee(nullptr, path_index);
+}

intern/cycles/kernel/device/optix/kernel_osl.cu

@@ -6,14 +6,14 @@
 
 /* Copy of the regular OptiX kernels with additional OSL support. */
 
-#include "kernel/device/optix/kernel_shader_raytrace.cu"
+#include "kernel/device/optix/kernel.cu"
 
 #include "kernel/bake/bake.h"
 #include "kernel/integrator/shade_background.h"
 #include "kernel/integrator/shade_dedicated_light.h"
 #include "kernel/integrator/shade_light.h"
 #include "kernel/integrator/shade_shadow.h"
-#include "kernel/integrator/shade_volume.h"
+#include "kernel/integrator/shade_surface.h"
 
 #include "kernel/device/gpu/work_stealing.h"
 
@@ -53,24 +53,6 @@ extern "C" __global__ void __raygen__kernel_optix_integrator_shade_surface()
   integrator_shade_surface(nullptr, path_index, kernel_params.render_buffer);
 }
 
-extern "C" __global__ void __raygen__kernel_optix_integrator_shade_volume()
-{
-  const int global_index = optixGetLaunchIndex().x;
-  const int path_index = (kernel_params.path_index_array) ?
-                             kernel_params.path_index_array[global_index] :
-                             global_index;
-  integrator_shade_volume(nullptr, path_index, kernel_params.render_buffer);
-}
-
-extern "C" __global__ void __raygen__kernel_optix_integrator_shade_volume_ray_marching()
-{
-  const int global_index = optixGetLaunchIndex().x;
-  const int path_index = (kernel_params.path_index_array) ?
-                             kernel_params.path_index_array[global_index] :
-                             global_index;
-  integrator_shade_volume_ray_marching(nullptr, path_index, kernel_params.render_buffer);
-}
-
 extern "C" __global__ void __raygen__kernel_optix_integrator_shade_shadow()
 {
   const int global_index = optixGetLaunchIndex().x;

intern/cycles/kernel/device/optix/kernel_osl_mnee.cu

@@ -0,0 +1,22 @@
+/* SPDX-FileCopyrightText: 2011-2026 Blender Foundation
+ *
+ * SPDX-License-Identifier: Apache-2.0 */
+
+#define WITH_OSL
+
+#include "kernel/device/optix/compat.h"
+#include "kernel/device/optix/globals.h"
+
+#include "kernel/bvh/bvh.h"
+#include "kernel/integrator/path_state.h"
+
+#include "kernel/integrator/intersect_mnee.h"
+
+extern "C" __global__ void __raygen__kernel_optix_integrator_intersect_mnee()
+{
+  const int global_index = optixGetLaunchIndex().x;
+  const int path_index = (kernel_params.path_index_array) ?
+                             kernel_params.path_index_array[global_index] :
+                             global_index;
+  integrator_intersect_mnee(nullptr, path_index);
+}

intern/cycles/kernel/device/optix/kernel_osl_shader_raytrace.cu

@@ -0,0 +1,19 @@
+/* SPDX-FileCopyrightText: 2011-2026 Blender Foundation
+ *
+ * SPDX-License-Identifier: Apache-2.0 */
+
+#define WITH_OSL
+
+#include "kernel/device/optix/compat.h"
+#include "kernel/device/optix/globals.h"
+
+#include "kernel/integrator/shade_surface.h"
+
+extern "C" __global__ void __raygen__kernel_optix_integrator_shade_surface_raytrace()
+{
+  const int global_index = optixGetLaunchIndex().x;
+  const int path_index = (kernel_params.path_index_array) ?
+                             kernel_params.path_index_array[global_index] :
+                             global_index;
+  integrator_shade_surface_raytrace(nullptr, path_index, kernel_params.render_buffer);
+}

intern/cycles/kernel/device/optix/kernel_osl_volume.cu

@@ -0,0 +1,34 @@
+/* SPDX-FileCopyrightText: 2011-2026 Blender Foundation
+ *
+ * SPDX-License-Identifier: Apache-2.0 */
+
+#define WITH_OSL
+
+/* Volume shading raygens for OSL, loaded as a separate optix module so they
+ * can be compiled in parallel with the base OSL module and skipped for
+ * scenes without volumes. */
+
+#include "kernel/device/optix/compat.h"
+#include "kernel/device/optix/globals.h"
+
+#include "kernel/film/data_passes.h"
+
+#include "kernel/integrator/shade_volume.h"
+
+extern "C" __global__ void __raygen__kernel_optix_integrator_shade_volume()
+{
+  const int global_index = optixGetLaunchIndex().x;
+  const int path_index = (kernel_params.path_index_array) ?
+                             kernel_params.path_index_array[global_index] :
+                             global_index;
+  integrator_shade_volume(nullptr, path_index, kernel_params.render_buffer);
+}
+
+extern "C" __global__ void __raygen__kernel_optix_integrator_shade_volume_ray_marching()
+{
+  const int global_index = optixGetLaunchIndex().x;
+  const int path_index = (kernel_params.path_index_array) ?
+                             kernel_params.path_index_array[global_index] :
+                             global_index;
+  integrator_shade_volume_ray_marching(nullptr, path_index, kernel_params.render_buffer);
+}

intern/cycles/kernel/device/optix/kernel_shader_raytrace.cu

@@ -1,13 +1,15 @@
-/* SPDX-FileCopyrightText: 2021-2022 Blender Foundation
+/* SPDX-FileCopyrightText: 2021-2026 Blender Foundation
  *
  * SPDX-License-Identifier: Apache-2.0 */
 
 /* Copy of the regular kernels with additional shader ray-tracing kernel that takes
  * much longer to compiler. This is only loaded when needed by the scene. */
 
-#include "kernel/device/optix/kernel.cu"
+#include "kernel/device/optix/compat.h"
+#include "kernel/device/optix/globals.h"
+
+#include "kernel/device/gpu/image.h"
 
-#include "kernel/integrator/intersect_mnee.h"
 #include "kernel/integrator/shade_surface.h"
 
 extern "C" __global__ void __raygen__kernel_optix_integrator_shade_surface_raytrace()
@@ -18,12 +20,3 @@ extern "C" __global__ void __raygen__kernel_optix_integrator_shade_surface_raytr
                              global_index;
   integrator_shade_surface_raytrace(nullptr, path_index, kernel_params.render_buffer);
 }
-
-extern "C" __global__ void __raygen__kernel_optix_integrator_intersect_mnee()
-{
-  const int global_index = optixGetLaunchIndex().x;
-  const int path_index = (kernel_params.path_index_array) ?
-                             kernel_params.path_index_array[global_index] :
-                             global_index;
-  integrator_intersect_mnee(nullptr, path_index);
-}

intern/cycles/kernel/integrator/state_util.h

@@ -14,6 +14,10 @@
 
 #include "kernel/util/differential.h"
 
+#if defined(__KERNEL_GPU__)
+#  include "util/atomic.h"
+#endif
+
 CCL_NAMESPACE_BEGIN
 
 /* Ray */