Add support for global shader parameters to OptiX path (#1323)

There are two main pieces here.

First, we specialize the code generaiton for CUDA kernels to account for the way that shader parameters are passed differently for ordinary compute kernels vs. ray-tracing kernels. Both global and entry-point shader parameters in Slang are translated to kernel function parameters for CUDA compute kernels, while for OptiX ray tracing kernels we need to use a global `__constant__` variable for the global parameters, and the SBT data (accessed via an OptiX API function) for entry-point shader parameters.

This choice bakes in a few pieces of policy when it comes to how Slang ray-tracing shaders translate to OptiX:

* It fixes the name used for the global `__constant__` variable for global shader parameters to be `SLANG_globalParams`. Since that name has to be specified when creating a pipeline with the OptiX API, the choice of name effectively becomes an ABI contract for Slang's code generation.

* It fixes the choice that global parameters in Slang map to per-launch parameters in OptiX, and entry-point parameters in Slang map to SBT-backed parameters in OptiX. This is a reasonable policy, and it is also one that we are likely to need to codify for Vulkan as well, but it is always a bit unfortunate to bake policy choices like this into the compiler (especially when shaders compiled for D3D can often decouple the form of their HLSL/Slang code from how things are bound in the API).

The second piece is a lot of refactoring of the logic in `render-test/cuda/cuda-compute-util.cpp`, so that the logic for setting up (and reading back) the buffers of parameter data can be shared between the compute and ray-tracing paths. The result may not be a true global optimum for how the code is organized, but it at least serves the goal of not duplicating the parameter-binding logic between compute and ray-tracing.

1 files changed, 79 insertions, 4 deletions

diff --git a/source/slang/slang-emit-cuda.cpp b/source/slang/slang-emit-cuda.cpp
index ac1e1ea63..fa63ba255 100644
--- a/source/slang/slang-emit-cuda.cpp
+++ b/source/slang/slang-emit-cuda.cpp
@@ -787,6 +787,22 @@ void CUDASourceEmitter::emitModuleImpl(IRModule* module)
             #undef CASE
                 }
 
+                if( stage != Stage::Compute )
+                {
+                    // Non-compute shaders (currently just OptiX ray tracing kernels)
+                    // require parameter data that is shared across multiple kernels
+                    // (which in our case is the global-scope shader parameters)
+                    // to be passed using a global `__constant__` variable.
+                    //
+                    // The use of `"C"` linkage here is required because the name
+                    // of this symbol must be passed to the OptiX API when creating
+                    // a pipeline that uses this compiled module. The exact name
+                    // used here (`SLANG_globalParams`) is thus a part of the
+                    // binary interface for Slang->OptiX translation.
+                    //
+                    m_writer->emit("extern \"C\" { __constant__ UniformState SLANG_globalParams; }\n");
+                }
+
                 // As a convenience for anybody reading the generated
                 // CUDA C++ code, we will prefix a compute kernel
                 // with the information from the `[numthreads(...)]`
@@ -815,27 +831,86 @@ void CUDASourceEmitter::emitModuleImpl(IRModule* module)
                 emitEntryPointAttributes(func, entryPointDecor);
                 emitType(resultType, globalSymbolName);
 
-                m_writer->emit("(UniformEntryPointParams* params, UniformState* uniformState)");
+                if( stage == Stage::Compute )
+                {
+                    // CUDA compute shaders take all of their parameters explicitly as
+                    // part of the entry-point parameter list. This means that the
+                    // data representing Slang shader parameters at both the global
+                    // and entry-point scopes needs to be passed as parameters.
+                    //
+                    // At the binary level, our generated CUDA compute kernels will take
+                    // two pointer parameters: the first points to the per-entry-point
+                    // `uniform` parameter data, and the second poinst to the global-scope
+                    // parameter data (if any).
+                    //
+                    m_writer->emit("(UniformEntryPointParams* entryPointShaderParameters, UniformState* uniformState)");
+                }
+                else
+                {
+                    // Non-compute shaders (currently just OptiX ray tracing kernels)
+                    // rely on other mechanisms for parameter passing, and thus use
+                    // an empty parameter list on the kernel declaration.
+                    //
+                    m_writer->emit("()");
+                }
+
                 emitSemantics(func);
                 m_writer->emit("\n{\n");
                 m_writer->indent();
 
                 // Initialize when constructing so that globals are zeroed
                 m_writer->emit("Context context = {};\n");
-                m_writer->emit("context.uniformState = uniformState;\n");
+
+                // The global-scope parameter data got passed in differently depending on whether we have
+                // a compute shader or a ray-tracing shader, so we need to alter how we initialize
+                // the pointer in our `context` based on the stage.
+                //
+                if( stage == Stage::Compute )
+                {
+                    m_writer->emit("context.uniformState = uniformState;\n");
+                }
+                else
+                {
+                    m_writer->emit("context.uniformState = &SLANG_globalParams;\n");
+                }
 
                 if (entryPointGlobalParams)
                 {
                     auto varDecl = entryPointGlobalParams;
                     auto rawType = varDecl->getDataType();
-
                     auto varType = rawType;
 
                     m_writer->emit("context.");
                     m_writer->emit(getName(varDecl));
                     m_writer->emit(" =  (");
                     emitType(varType);
-                    m_writer->emit("*)params; \n");
+                    m_writer->emit("*)");
+
+                    // Similar to the case for global parameter data above, the entry-point
+                    // uniform parameter data gets passed in differently for compute kernels
+                    // vs. ray-tracing kernels, and we need to handle the two cases here.
+                    //
+                    if( stage == Stage::Compute )
+                    {
+                        // In the compute case, the entry-point uniform parameters came
+                        // in as an explicit parameter on the CUDA kernel, and we simply
+                        // cast it to the expected type here.
+                        //
+                        m_writer->emit("entryPointShaderParameters");
+                    }
+                    else
+                    {
+                        // In the ray-tracing case, the entry-point uniform parameters
+                        // implicitly map to the contents of the Shader Binding Table
+                        // (SBT) entry for the entry point instance being invoked.
+                        //
+                        // The OptiX API provides an accessor function to get a pointer
+                        // to the SBT data for the current entry, and we cast the result
+                        // of that to the expected type.
+                        //
+                        m_writer->emit("optixGetSbtDataPointer()");
+                    }
+                    m_writer->emit(";\n");
                 }
 
                 m_writer->emit("context.");