Cycles: Move OptiX OSL Camera kernel into its own PTX module

On the one hand, this improves initialization time since we don't need to
load/compile the full OSL module with all the shading logic if we're only
using a custom camera with SVM shading.

On the other hand, it also fixes a bug I noticed while preparing test scenes:
The AO and Bevel nodes don't work when using custom cameras with SVM on OptiX.

The issue there is that those two are handled by the SHADE_SURFACE_RAYTRACE
kernel, but since that one has intersection logic, we use the OptiX-specific
kernel even if OSL shading is disabled.
However, with the previous unified OSL module, this would mean loading
SHADE_SURFACE_RAYTRACE from kernel_osl.cu, which has `#undef __SVM__` and
therefore doesn't handle them correctly.

With this change, we'll use the kernels from kernel_shader_raytrace.cu in that
case, which do support SVM nodes just fine.

Disk usage of the new kernel_optix_osl_camera.ptx.zst file is 30KB, so this
also doesn't blow up the kernel disk size (and kernel_optix_osl.ptx.zst is
probably smaller by that amount now).

Since it seems that we can mix modules just fine, I'm suspecting that we could
split the modules properly (intersection, SVM shading with raytracing,
OSL shading, OSL camera), instead of the current approach where modules
essentially correspond to feature set tiers and each includes the previous
one's kernels as well - but that's a separate refactor.

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

intern/cycles/device/optix/device_impl.cpp

@@ -132,6 +132,9 @@ OptiXDevice::~OptiXDevice()
   }
 
 #  ifdef WITH_OSL
+  if (osl_camera_module != nullptr) {
+    optixModuleDestroy(osl_camera_module);
+  }
   for (const OptixModule &module : osl_modules) {
     if (module != nullptr) {
       optixModuleDestroy(module);
@@ -199,10 +202,11 @@ 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 uint osl_mask = KERNEL_FEATURE_OSL_SHADING | KERNEL_FEATURE_OSL_CAMERA;
-  const bool use_osl = (kernel_features & osl_mask);
+  const bool use_osl_shading = (kernel_features & KERNEL_FEATURE_OSL_SHADING);
+  const bool use_osl_camera = (kernel_features & KERNEL_FEATURE_OSL_CAMERA);
 #  else
-  const bool use_osl = false;
+  const bool use_osl_shading = false;
+  const bool use_osl_camera = false;
 #  endif
 
   /* Skip creating OptiX module if only doing denoising. */
@@ -212,10 +216,10 @@ 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 ? "_osl" :
+  string suffix = use_osl_shading ? "_osl" :
                   (kernel_features & (KERNEL_FEATURE_NODE_RAYTRACE | KERNEL_FEATURE_MNEE)) ?
-                            "_shader_raytrace" :
-                            "";
+                                    "_shader_raytrace" :
+                                    "";
   string ptx_filename;
   if (need_optix_kernels) {
     ptx_filename = path_get("lib/kernel_optix" + suffix + ".ptx.zst");
@@ -274,6 +278,11 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   }
 
 #  ifdef WITH_OSL
+  if (osl_camera_module != nullptr) {
+    optixModuleDestroy(osl_camera_module);
+    osl_camera_module = nullptr;
+  }
+
   /* Recreating base OptiX module invalidates all OSL modules too, since they link against it. */
   for (const OptixModule &module : osl_modules) {
     if (module != nullptr) {
@@ -515,8 +524,9 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
     group_descs[PG_RGEN_SHADE_SURFACE_RAYTRACE].raygen.entryFunctionName =
         "__raygen__kernel_optix_integrator_shade_surface_raytrace";
 
-    /* Kernels with OSL support are built without SVM, so can skip those direct callables there. */
-    if (!use_osl) {
+    /* Kernels with OSL shading support are built without SVM, so can skip those direct callables
+     * 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.entryFunctionNameDC = "__direct_callable__svm_node_ao";
@@ -535,7 +545,7 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   }
 
   /* OSL uses direct callables to execute, so shading needs to be done in OptiX if OSL is used. */
-  if (use_osl) {
+  if (use_osl_shading) {
     group_descs[PG_RGEN_SHADE_BACKGROUND].kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
     group_descs[PG_RGEN_SHADE_BACKGROUND].raygen.module = optix_module;
     group_descs[PG_RGEN_SHADE_BACKGROUND].raygen.entryFunctionName =
@@ -572,11 +582,51 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
     group_descs[PG_RGEN_EVAL_CURVE_SHADOW_TRANSPARENCY].raygen.module = optix_module;
     group_descs[PG_RGEN_EVAL_CURVE_SHADOW_TRANSPARENCY].raygen.entryFunctionName =
         "__raygen__kernel_optix_shader_eval_curve_shadow_transparency";
+  }
+
+#  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;
+    }
+
+#    if OPTIX_ABI_VERSION >= 84
+    const OptixResult result = optixModuleCreate(context,
+                                                 &module_options,
+                                                 &pipeline_options,
+                                                 ptx_data.data(),
+                                                 ptx_data.size(),
+                                                 nullptr,
+                                                 nullptr,
+                                                 &osl_camera_module);
+#    else
+    const OptixResult result = optixModuleCreateFromPTX(context,
+                                                        &module_options,
+                                                        &pipeline_options,
+                                                        ptx_data.data(),
+                                                        ptx_data.size(),
+                                                        nullptr,
+                                                        nullptr,
+                                                        &osl_camera_module);
+#    endif
+    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 = optix_module;
+    group_descs[PG_RGEN_INIT_FROM_CAMERA].raygen.module = osl_camera_module;
     group_descs[PG_RGEN_INIT_FROM_CAMERA].raygen.entryFunctionName =
         "__raygen__kernel_optix_integrator_init_from_camera";
   }
+#  endif
 
   optix_assert(optixProgramGroupCreate(
       context, group_descs, NUM_PROGRAM_GROUPS, &group_options, nullptr, nullptr, groups));
@@ -618,7 +668,7 @@ bool OptiXDevice::load_kernels(const uint kernel_features)
   link_options.debugLevel = module_options.debugLevel;
 #  endif
 
-  if (use_osl) {
+  if (use_osl_shading || use_osl_camera) {
     /* OSL kernels will be (re)created on by OSL manager. */
   }
   else if (kernel_features & (KERNEL_FEATURE_NODE_RAYTRACE | KERNEL_FEATURE_MNEE)) {
@@ -952,6 +1002,8 @@ bool OptiXDevice::load_osl_kernels()
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_LIGHT]);
     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_SHADE_SURFACE_MNEE]);
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_VOLUME]);
     pipeline_groups.push_back(groups[PG_RGEN_SHADE_SHADOW]);
@@ -1000,7 +1052,8 @@ bool OptiXDevice::load_osl_kernels()
 
     const unsigned int css = std::max(stack_size[PG_RGEN_SHADE_SURFACE_RAYTRACE].cssRG,
                                       stack_size[PG_RGEN_SHADE_SURFACE_MNEE].cssRG);
-    unsigned int dss = 0;
+    unsigned int dss = std::max(stack_size[PG_CALL_SVM_AO].dssDC,
+                                stack_size[PG_CALL_SVM_BEVEL].dssDC);
     for (unsigned int i = 0; i < osl_stack_size.size(); ++i) {
       dss = std::max(dss, osl_stack_size[i].dssDC);
     }

intern/cycles/device/optix/device_impl.h

@@ -81,6 +81,7 @@ class OptiXDevice : public CUDADevice {
   OSLGlobals osl_globals;
   vector<OptixModule> osl_modules;
   vector<OptixProgramGroup> osl_groups;
+  OptixModule osl_camera_module = nullptr;
 #  endif
 
  private:

intern/cycles/kernel/CMakeLists.txt

@@ -46,6 +46,7 @@ if(WITH_CYCLES_OSL)
     ${SRC_KERNEL_DEVICE_OPTIX}
     osl/services_optix.cu
     device/optix/kernel_osl.cu
+    device/optix/kernel_osl_camera.cu
   )
 endif()
 
@@ -916,6 +917,10 @@ if(WITH_CYCLES_DEVICE_OPTIX AND WITH_CYCLES_CUDA_BINARIES)
       kernel_optix_osl
       "device/optix/kernel_osl.cu"
       "--relocatable-device-code=true")
+    cycles_optix_kernel_add(
+      kernel_optix_osl_camera
+      "device/optix/kernel_osl_camera.cu"
+      "--relocatable-device-code=true")
     cycles_optix_kernel_add(
       kernel_optix_osl_services
       "osl/services_optix.cu"

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

@@ -9,7 +9,6 @@
 #include "kernel/device/optix/kernel_shader_raytrace.cu"
 
 #include "kernel/bake/bake.h"
-#include "kernel/integrator/init_from_camera.h"
 #include "kernel/integrator/shade_background.h"
 #include "kernel/integrator/shade_dedicated_light.h"
 #include "kernel/integrator/shade_light.h"
@@ -95,26 +94,3 @@ extern "C" __global__ void __raygen__kernel_optix_shader_eval_curve_shadow_trans
   const int global_index = kernel_params.offset + optixGetLaunchIndex().x;
   kernel_curve_shadow_transparency_evaluate(nullptr, input, output, global_index);
 }
-
-extern "C" __global__ void __raygen__kernel_optix_integrator_init_from_camera()
-{
-  const int global_index = optixGetLaunchIndex().x;
-
-  const KernelWorkTile *tiles = (const KernelWorkTile *)kernel_params.path_index_array;
-
-  const int tile_index = global_index / kernel_params.max_tile_work_size;
-  const int tile_work_index = global_index - tile_index * kernel_params.max_tile_work_size;
-
-  const KernelWorkTile *tile = &tiles[tile_index];
-
-  if (tile_work_index >= tile->work_size) {
-    return;
-  }
-
-  const int path_index = tile->path_index_offset + tile_work_index;
-
-  uint x, y, sample;
-  get_work_pixel(tile, tile_work_index, &x, &y, &sample);
-
-  integrator_init_from_camera(nullptr, path_index, tile, kernel_params.render_buffer, x, y, sample);
-}

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

@@ -0,0 +1,35 @@
+/* SPDX-FileCopyrightText: 2011-2025 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/init_from_camera.h"
+
+#include "kernel/device/gpu/work_stealing.h"
+
+extern "C" __global__ void __raygen__kernel_optix_integrator_init_from_camera()
+{
+  const int global_index = optixGetLaunchIndex().x;
+
+  const KernelWorkTile *tiles = (const KernelWorkTile *)kernel_params.path_index_array;
+
+  const int tile_index = global_index / kernel_params.max_tile_work_size;
+  const int tile_work_index = global_index - tile_index * kernel_params.max_tile_work_size;
+
+  const KernelWorkTile *tile = &tiles[tile_index];
+
+  if (tile_work_index >= tile->work_size) {
+    return;
+  }
+
+  const int path_index = tile->path_index_offset + tile_work_index;
+
+  uint x, y, sample;
+  get_work_pixel(tile, tile_work_index, &x, &y, &sample);
+
+  integrator_init_from_camera(nullptr, path_index, tile, kernel_params.render_buffer, x, y, sample);
+}

intern/cycles/kernel/film/write.h

@@ -10,6 +10,8 @@
 
 #include "kernel/util/colorspace.h"
 
+#include "util/atomic.h"
+
 #ifdef __KERNEL_GPU__
 #  define __ATOMIC_PASS_WRITE__
 #endif