Initial framework for compute shaders on WebGPU

src/core/constants.js

@@ -104,6 +104,13 @@ export const TRACEID_BINDGROUPFORMAT_ALLOC = 'BindGroupFormatAlloc';
  */
 export const TRACEID_RENDERPIPELINE_ALLOC = 'RenderPipelineAlloc';
 
+/**
+ * Logs the creation of compute pipelines. WebGPU only.
+ *
+ * @type {string}
+ */
+export const TRACEID_COMPUTEPIPELINE_ALLOC = 'ComputePipelineAlloc';
+
 /**
  * Logs the creation of pipeline layouts. WebBPU only.
  *

src/core/tracing.js

@@ -37,6 +37,7 @@ class Tracing {
      * - {@link TRACEID_VRAM_VB}
      * - {@link TRACEID_VRAM_IB}
      * - {@link TRACEID_RENDERPIPELINE_ALLOC}
+     * - {@link TRACEID_COMPUTEPIPELINE_ALLOC}
      * - {@link TRACEID_PIPELINELAYOUT_ALLOC}
      * - {@link TRACEID_TEXTURES}
      * - {@link TRACEID_GPU_TIMINGS}

src/index.js

@@ -58,6 +58,7 @@ export * from './platform/audio/constants.js';
 export * from './platform/graphics/constants.js';
 export { createGraphicsDevice } from './platform/graphics/graphics-device-create.js';
 export { BlendState } from './platform/graphics/blend-state.js';
+export { Compute } from './platform/graphics/compute.js';
 export { DepthState } from './platform/graphics/depth-state.js';
 export { GraphicsDevice } from './platform/graphics/graphics-device.js';
 export { IndexBuffer } from './platform/graphics/index-buffer.js';

src/platform/graphics/compute.js

@@ -0,0 +1,44 @@
+/**
+ * A representation of a compute shader with the associated data, that can be executed on the GPU.
+ *
+ * @ignore
+ */
+class Compute {
+    /**
+     * A compute shader.
+     *
+     * @type {import('./shader.js').Shader|null}
+     * @ignore
+     */
+    shader = null;
+
+    /**
+     * Create a compute instance. Note that this is supported on WebGPU only and is a no-op on
+     * other platforms.
+     *
+     * @param {import('./graphics-device.js').GraphicsDevice} graphicsDevice -
+     * The graphics device.
+     * @param {import('./shader.js').Shader} shader - The compute shader.
+     */
+    constructor(graphicsDevice, shader) {
+        this.device = graphicsDevice;
+        this.shader = shader;
+
+        if (graphicsDevice.supportsCompute) {
+            this.impl = graphicsDevice.createComputeImpl(this);
+        }
+    }
+
+    /**
+     * Dispatch the compute work.
+     *
+     * @param {number} x - X dimension of the grid of work-groups to dispatch.
+     * @param {number} [y] - Y dimension of the grid of work-groups to dispatch.
+     * @param {number} [z] - Z dimension of the grid of work-groups to dispatch.
+     */
+    dispatch(x, y, z) {
+        this.impl?.dispatch(x, y, z);
+    }
+}
+
+export { Compute };

src/platform/graphics/graphics-device.js

@@ -187,6 +187,14 @@ class GraphicsDevice extends EventHandler {
      */
     supportsVolumeTextures = false;
 
+    /**
+     * True if the device supports compute shaders.
+     *
+     * @readonly
+     * @type {boolean}
+     */
+    supportsCompute = false;
+
     /**
      * Currently active render target.
      *
@@ -753,6 +761,18 @@ class GraphicsDevice extends EventHandler {
         this.boneLimit = maxBones;
     }
 
+    startRenderPass(renderPass) {
+    }
+
+    endRenderPass(renderPass) {
+    }
+
+    startComputePass() {
+    }
+
+    endComputePass() {
+    }
+
     /**
      * Function which executes at the start of the frame. This should not be called manually, as
      * it is handled by the AppBase instance.

src/platform/graphics/null/null-graphics-device.js

@@ -123,12 +123,6 @@ class NullGraphicsDevice extends GraphicsDevice {
         super.initializeContextCaches();
     }
 
-    startPass(renderPass) {
-    }
-
-    endPass(renderPass) {
-    }
-
     clear(options) {
     }
 

src/platform/graphics/render-pass.js

@@ -286,13 +286,13 @@ class RenderPass {
             if (this.executeEnabled) {
 
                 if (realPass) {
-                    device.startPass(this);
+                    device.startRenderPass(this);
                 }
 
                 this.execute();
 
                 if (realPass) {
-                    device.endPass(this);
+                    device.endRenderPass(this);
                 }
             }
 

src/platform/graphics/shader.js

@@ -46,9 +46,12 @@ class Shader {
      * vertex shader attribute names to semantics SEMANTIC_*. This enables the engine to match
      * vertex buffer data as inputs to the shader. When not specified, rendering without
      * vertex buffer is assumed.
-     * @param {string} definition.vshader - Vertex shader source (GLSL code).
+     * @param {string} [definition.vshader] - Vertex shader source (GLSL code). Optional when
+     * compute shader is specified.
      * @param {string} [definition.fshader] - Fragment shader source (GLSL code). Optional when
-     * useTransformFeedback is specified.
+     * useTransformFeedback or compute shader is specified.
+     * @param {string} [definition.cshader] - Compute shader source (WGSL code). Only supported on
+     * WebGPU platform.
      * @param {boolean} [definition.useTransformFeedback] - Specifies that this shader outputs
      * post-VS data to a buffer.
      * @param {string} [definition.shaderLanguage] - Specifies the shader language of vertex and
@@ -89,15 +92,19 @@ class Shader {
         this.device = graphicsDevice;
         this.definition = definition;
         this.name = definition.name || 'Untitled';
+        this.init();
 
-        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.');
-
-        // pre-process shader sources
-        definition.vshader = Preprocessor.run(definition.vshader);
-        definition.fshader = Preprocessor.run(definition.fshader, graphicsDevice.isWebGL2);
+        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.');
+        } 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.');
 
-        this.init();
+            // pre-process shader sources
+            definition.vshader = Preprocessor.run(definition.vshader);
+            definition.fshader = Preprocessor.run(definition.fshader, graphicsDevice.isWebGL2);
+        }
 
         this.impl = graphicsDevice.createShaderImpl(this);
 

src/platform/graphics/webgl/webgl-graphics-device.js

@@ -1408,7 +1408,7 @@ class WebglGraphicsDevice extends GraphicsDevice {
      * @param {import('../render-pass.js').RenderPass} renderPass - The render pass to start.
      * @ignore
      */
-    startPass(renderPass) {
+    startRenderPass(renderPass) {
 
         DebugGraphics.pushGpuMarker(this, `START-PASS`);
 
@@ -1470,7 +1470,7 @@ class WebglGraphicsDevice extends GraphicsDevice {
      * @param {import('../render-pass.js').RenderPass} renderPass - The render pass to end.
      * @ignore
      */
-    endPass(renderPass) {
+    endRenderPass(renderPass) {
 
         DebugGraphics.pushGpuMarker(this, `END-PASS`);
 

src/platform/graphics/webgpu/webgpu-compute-pipeline.js

@@ -0,0 +1,61 @@
+import { Debug, DebugHelper } from "../../../core/debug.js";
+import { TRACEID_COMPUTEPIPELINE_ALLOC } from "../../../core/constants.js";
+import { WebgpuDebug } from "./webgpu-debug.js";
+
+let _pipelineId = 0;
+
+/**
+ * @ignore
+ */
+class WebgpuComputePipeline {
+    constructor(device) {
+        /** @type {import('./webgpu-graphics-device.js').WebgpuGraphicsDevice} */
+        this.device = device;
+    }
+
+    /** @private */
+    get(shader) {
+
+        const pipeline = this.create(shader);
+        return pipeline;
+    }
+
+    create(shader) {
+
+        const wgpu = this.device.wgpu;
+
+        /** @type {import('./webgpu-shader.js').WebgpuShader} */
+        const webgpuShader = shader.impl;
+
+        /** @type {GPUComputePipelineDescriptor} */
+        const descr = {
+            compute: {
+                module: webgpuShader.getComputeShaderModule(),
+                entryPoint: webgpuShader.computeEntryPoint
+            },
+
+            // uniform / texture binding layout
+            layout: 'auto'
+        };
+
+        WebgpuDebug.validate(this.device);
+
+        _pipelineId++;
+        DebugHelper.setLabel(descr, `ComputePipelineDescr-${_pipelineId}`);
+
+        const pipeline = wgpu.createComputePipeline(descr);
+
+        DebugHelper.setLabel(pipeline, `ComputePipeline-${_pipelineId}`);
+        Debug.trace(TRACEID_COMPUTEPIPELINE_ALLOC, `Alloc: Id ${_pipelineId}`, descr);
+
+        WebgpuDebug.end(this.device, {
+            computePipeline: this,
+            descr,
+            shader
+        });
+
+        return pipeline;
+    }
+}
+
+export { WebgpuComputePipeline };

src/platform/graphics/webgpu/webgpu-compute.js

@@ -0,0 +1,30 @@
+/**
+ * A WebGPU implementation of the Compute.
+ *
+ * @ignore
+ */
+class WebgpuCompute {
+    constructor(compute) {
+        this.compute = compute;
+
+        const { device, shader } = compute;
+        this.pipeline = device.computePipeline.get(shader);
+    }
+
+    dispatch(x, y, z) {
+
+        // TODO: currently each dispatch is a separate compute pass, which is not optimal, and we should
+        // batch multiple dispatches into a single compute pass
+        const device = this.compute.device;
+        device.startComputePass();
+
+        // dispatch
+        const passEncoder = device.passEncoder;
+        passEncoder.setPipeline(this.pipeline);
+        passEncoder.dispatchWorkgroups(x, y, z);
+
+        device.endComputePass();
+    }
+}
+
+export { WebgpuCompute };