Compute shaders preprocess the provided source code (#7783)

Co-authored-by: Martin Valigursky <[email protected]>

Author: mvaligursky

examples/src/examples/compute/indirect-draw.compute-shader.wgsl

@@ -1,11 +1,6 @@
-// Indexed draw call parameters
-struct DrawIndexedIndirectArgs {
-    indexCount: u32,
-    instanceCount: u32,
-    firstIndex: u32,
-    baseVertex: i32,
-    firstInstance: u32
-};
+// include built-in engine chunk, which gives us DrawIndexedIndirectArgs (but also DrawIndirectArgs)
+// structs, defining the format of the indirect draw parameters / buffer
+#include "indirectCoreCS"
 
 // Binding 0: uniform buffer holding draw call metadata and runtime config
 struct Uniforms {

examples/src/examples/compute/indirect-draw.example.mjs

@@ -121,6 +121,9 @@ assetListLoader.load(() => {
             shaderLanguage: pc.SHADERLANGUAGE_WGSL,
             cshader: files['compute-shader.wgsl'],
 
+            // include all WGSL chunks to be available for including in the compute shader
+            cincludes: pc.ShaderChunks.get(device, pc.SHADERLANGUAGE_WGSL),
+
             // format of a uniform buffer used by the compute shader
             computeUniformBufferFormats: {
                 ub: new pc.UniformBufferFormat(device, [

src/platform/graphics/shader.js

@@ -75,6 +75,9 @@ class Shader {
      * @param {Map<string, string>} [definition.fincludes] - A map containing key-value pairs
      * of include names and their content. These are used for resolving #include directives in the
      * fragment shader source.
+     * @param {Map<string, string>} [definition.cincludes] - A map containing key-value pairs
+     * of include names and their content. These are used for resolving #include directives in the
+     * compute shader source.
      * @param {boolean} [definition.useTransformFeedback] - Specifies that this shader outputs
      * post-VS data to a buffer.
      * @param {string | string[]} [definition.fragmentOutputTypes] - Fragment shader output types,
@@ -123,6 +126,13 @@ class Shader {
         if (definition.cshader) {
             Debug.assert(graphicsDevice.supportsCompute, 'Compute shaders are not supported on this device.');
             Debug.assert(!definition.vshader && !definition.fshader, 'Vertex and fragment shaders are not supported when creating a compute shader.');
+
+            // pre-process compute shader source
+            definition.cshader = Preprocessor.run(definition.cshader, definition.cincludes, {
+                sourceName: `compute shader for ${this.label}`,
+                stripDefines: true
+            });
+
         } else {
             Debug.assert(definition.vshader, 'No vertex shader has been specified when creating a shader.');
             Debug.assert(definition.fshader, 'No fragment shader has been specified when creating a shader.');

src/scene/shader-lib/wgsl/chunks/common/comp/indirect-core.js

@@ -0,0 +1,19 @@
+export default /* wgsl */`
+// Indirect indexed draw call parameters
+struct DrawIndexedIndirectArgs {
+    indexCount: u32,        // Number of indices to draw
+    instanceCount: u32,     // Number of instances to draw
+    firstIndex: u32,        // Index of the first index in the index buffer
+    baseVertex: i32,        // Offset added to each index before fetching vertex
+    firstInstance: u32      // First instance ID
+};
+
+// Indirect non-indexed draw call parameters
+struct DrawIndirectArgs {
+    vertexCount: u32,       // Number of vertices to draw
+    instanceCount: u32,     // Number of instances to draw
+    firstVertex: u32,       // Index of the first vertex
+    firstInstance: u32,     // First instance ID
+    _pad: u32               // Padding to match indexed size (unused)
+};
+`;

src/scene/shader-lib/wgsl/collections/shader-chunks-wgsl.js

@@ -59,6 +59,7 @@ import gsplatSHVS from '../chunks/gsplat/vert/gsplatSH.js';
 import gsplatSourceVS from '../chunks/gsplat/vert/gsplatSource.js';
 import gsplatVS from '../chunks/gsplat/vert/gsplat.js';
 import quadVS from '../chunks/common/vert/quad.js';
+import indirectCoreCS from '../chunks/common/comp/indirect-core.js';
 import immediateLinePS from '../chunks/internal/frag/immediateLine.js';
 import immediateLineVS from '../chunks/internal/vert/immediateLine.js';
 import iridescenceDiffractionPS from '../chunks/lit/frag/iridescenceDiffraction.js';
@@ -234,6 +235,7 @@ const shaderChunksWGSL = {
     gsplatSourceVS,
     gsplatVS,
     quadVS,
+    indirectCoreCS,
     immediateLinePS,
     immediateLineVS,
     iridescenceDiffractionPS,