Add compute dispatch indirect API

packages/dev/core/src/Compute/computeShader.ts

@@ -348,62 +348,81 @@ export class ComputeShader {
      * @returns True if the dispatch could be done, else false (meaning either the compute effect or at least one of the bound resources was not ready)
      */
     public dispatch(x: number, y?: number, z?: number): boolean {
-        if (!this.fastMode) {
-            if (!this.isReady()) {
-                return false;
-            }
+        if (!this.fastMode && !this._checkContext()) {
+            return false;
+        }
+        this._engine.computeDispatch(this._effect, this._context, this._bindings, x, y, z, this._options.bindingsMapping, this.gpuTimeInFrame);
 
-            // If the sampling parameters of a texture bound to the shader have changed, we must clear the compute context so that it is recreated with the updated values
-            // Also, if the actual (gpu) buffer used by a uniform buffer has changed, we must clear the compute context so that it is recreated with the updated value
-            for (const key in this._bindings) {
-                const binding = this._bindings[key];
+        return true;
+    }
 
-                if (!this._options.bindingsMapping[key]) {
-                    throw new Error("ComputeShader ('" + this.name + "'): No binding mapping has been provided for the property '" + key + "'");
-                }
+    /**
+     * Dispatches (executes) the compute shader.
+     * @param buffer Buffer containing the number of workgroups to execute on the X, Y and Z dimensions
+     * @param offset Offset in the buffer where the workgroup counts are stored (default: 0)
+     * @returns True if the dispatch could be done, else false (meaning either the compute effect or at least one of the bound resources was not ready)
+     */
+    public dispatchIndirect(buffer: StorageBuffer | DataBuffer, offset: number = 0): boolean {
+        if (!this.fastMode && !this._checkContext()) {
+            return false;
+        }
+        const dataBuffer = ComputeShader._BufferIsDataBuffer(buffer) ? buffer : buffer.getBuffer();
+        this._engine.computeDispatchIndirect(this._effect, this._context, this._bindings, dataBuffer, offset, this._options.bindingsMapping, this.gpuTimeInFrame);
+        return true;
+    }
 
-                switch (binding.type) {
-                    case ComputeBindingType.Texture: {
-                        const sampler = this._samplers[key];
-                        const texture = binding.object as BaseTexture;
-
-                        if (!sampler || !texture._texture || !sampler.compareSampler(texture._texture)) {
-                            this._samplers[key] = new TextureSampler().setParameters(
-                                texture.wrapU,
-                                texture.wrapV,
-                                texture.wrapR,
-                                texture.anisotropicFilteringLevel,
-                                texture._texture!.samplingMode,
-                                texture._texture?._comparisonFunction
-                            );
-                            this._contextIsDirty = true;
-                        }
-                        break;
-                    }
-                    case ComputeBindingType.ExternalTexture: {
-                        // we must recreate the bind groups each time if there's an external texture, because device.importExternalTexture must be called each frame
+    private _checkContext(): boolean {
+        if (!this.isReady()) {
+            return false;
+        }
+
+        // If the sampling parameters of a texture bound to the shader have changed, we must clear the compute context so that it is recreated with the updated values
+        // Also, if the actual (gpu) buffer used by a uniform buffer has changed, we must clear the compute context so that it is recreated with the updated value
+        for (const key in this._bindings) {
+            const binding = this._bindings[key];
+
+            if (!this._options.bindingsMapping[key]) {
+                throw new Error("ComputeShader ('" + this.name + "'): No binding mapping has been provided for the property '" + key + "'");
+            }
+
+            switch (binding.type) {
+                case ComputeBindingType.Texture: {
+                    const sampler = this._samplers[key];
+                    const texture = binding.object as BaseTexture;
+
+                    if (!sampler || !texture._texture || !sampler.compareSampler(texture._texture)) {
+                        this._samplers[key] = new TextureSampler().setParameters(
+                            texture.wrapU,
+                            texture.wrapV,
+                            texture.wrapR,
+                            texture.anisotropicFilteringLevel,
+                            texture._texture!.samplingMode,
+                            texture._texture?._comparisonFunction
+                        );
                         this._contextIsDirty = true;
-                        break;
                     }
-                    case ComputeBindingType.UniformBuffer: {
-                        const ubo = binding.object as UniformBuffer;
-                        if (ubo.getBuffer() !== binding.buffer) {
-                            binding.buffer = ubo.getBuffer();
-                            this._contextIsDirty = true;
-                        }
-                        break;
+                    break;
+                }
+                case ComputeBindingType.ExternalTexture: {
+                    // we must recreate the bind groups each time if there's an external texture, because device.importExternalTexture must be called each frame
+                    this._contextIsDirty = true;
+                    break;
+                }
+                case ComputeBindingType.UniformBuffer: {
+                    const ubo = binding.object as UniformBuffer;
+                    if (ubo.getBuffer() !== binding.buffer) {
+                        binding.buffer = ubo.getBuffer();
+                        this._contextIsDirty = true;
                     }
+                    break;
                 }
             }
-
-            if (this._contextIsDirty) {
-                this._contextIsDirty = false;
-                this._context.clear();
-            }
         }
 
-        this._engine.computeDispatch(this._effect, this._context, this._bindings, x, y, z, this._options.bindingsMapping, this.gpuTimeInFrame);
-
+        if (this._contextIsDirty) {
+            this._contextIsDirty = false;
+            this._context.clear();
+        }
         return true;
     }
 

packages/dev/core/src/Engines/Extensions/engine.computeShader.ts

@@ -6,6 +6,7 @@ import { ThinEngine } from "../../Engines/thinEngine";
 import type { Nullable } from "../../types";
 import { AbstractEngine } from "../abstractEngine";
 import type { WebGPUPerfCounter } from "../WebGPU/webgpuPerfCounter";
+import type { DataBuffer } from "../../Buffers/dataBuffer";
 
 /**
  * Type used to locate a resource in a compute shader.
@@ -114,6 +115,27 @@ declare module "../../Engines/abstractEngine" {
             gpuPerfCounter?: WebGPUPerfCounter
         ): void;
 
+        /**
+         * Dispatches a compute shader
+         * @param effect The compute effect
+         * @param context The compute context
+         * @param bindings The list of resources to bind to the shader
+         * @param x The number of workgroups to execute on the X dimension
+         * @param y The number of workgroups to execute on the Y dimension
+         * @param z The number of workgroups to execute on the Z dimension
+         * @param bindingsMapping list of bindings mapping (key is property name, value is binding location)
+         * @param gpuPerfCounter GPU time computed for the compute shader will be assigned to this object
+         */
+        computeDispatchIndirect(
+            effect: ComputeEffect,
+            context: IComputeContext,
+            bindings: ComputeBindingList,
+            buffer: DataBuffer,
+            offset?: number,
+            bindingsMapping?: ComputeBindingMapping,
+            gpuPerfCounter?: WebGPUPerfCounter
+        ): void;
+
         /**
          * Gets a boolean indicating if all created compute effects are ready
          * @returns true if all effects are ready
@@ -171,6 +193,16 @@ ThinEngine.prototype.computeDispatch = function (
 ): void {
     throw new Error("computeDispatch: This engine does not support compute shaders!");
 };
+ThinEngine.prototype.computeDispatchIndirect = function (
+    effect: ComputeEffect,
+    context: IComputeContext,
+    bindings: ComputeBindingList,
+    buffer: DataBuffer,
+    offset?: number,
+    bindingsMapping?: ComputeBindingMapping
+): void {
+    throw new Error("computeDispatchIndirect: This engine does not support compute shaders!");
+};
 
 ThinEngine.prototype.areAllComputeEffectsReady = function (): boolean {
     return true;

packages/dev/core/src/Engines/WebGPU/Extensions/engine.computeShader.ts

@@ -10,11 +10,27 @@ import { WebGPUComputeContext } from "../webgpuComputeContext";
 import { WebGPUComputePipelineContext } from "../webgpuComputePipelineContext";
 import * as WebGPUConstants from "../webgpuConstants";
 import type { WebGPUPerfCounter } from "../webgpuPerfCounter";
+import type { DataBuffer } from "../../../Buffers/dataBuffer";
 
 declare module "../../webgpuEngine" {
     export interface WebGPUEngine {
         /** @internal */
         _createComputePipelineStageDescriptor(computeShader: string, defines: Nullable<string>, entryPoint: string): GPUProgrammableStage;
+        /** @internal
+         * Either all of x,y,z or buffer and offset should be defined.
+         */
+        _computeDispatch(
+            effect: ComputeEffect,
+            context: IComputeContext,
+            bindings: ComputeBindingList,
+            x?: number,
+            y?: number,
+            z?: number,
+            buffer?: DataBuffer,
+            offset?: number,
+            bindingsMapping?: ComputeBindingMapping,
+            gpuPerfCounter?: WebGPUPerfCounter
+        ): void;
     }
 }
 
@@ -67,6 +83,33 @@ WebGPUEngine.prototype.computeDispatch = function (
     z = 1,
     bindingsMapping?: ComputeBindingMapping,
     gpuPerfCounter?: WebGPUPerfCounter
+): void {
+    this._computeDispatch(effect, context, bindings, x, y, z, undefined, undefined, bindingsMapping, gpuPerfCounter);
+};
+
+WebGPUEngine.prototype.computeDispatchIndirect = function (
+    effect: ComputeEffect,
+    context: IComputeContext,
+    bindings: ComputeBindingList,
+    buffer: DataBuffer,
+    offset: number = 0,
+    bindingsMapping?: ComputeBindingMapping,
+    gpuPerfCounter?: WebGPUPerfCounter
+): void {
+    this._computeDispatch(effect, context, bindings, undefined, undefined, undefined, buffer, offset, bindingsMapping, gpuPerfCounter);
+};
+
+WebGPUEngine.prototype._computeDispatch = function (
+    effect: ComputeEffect,
+    context: IComputeContext,
+    bindings: ComputeBindingList,
+    x?: number,
+    y?: number,
+    z?: number,
+    buffer?: DataBuffer,
+    offset?: number,
+    bindingsMapping?: ComputeBindingMapping,
+    gpuPerfCounter?: WebGPUPerfCounter
 ): void {
     this._endCurrentRenderPass();
 
@@ -97,8 +140,12 @@ WebGPUEngine.prototype.computeDispatch = function (
         computePass.setBindGroup(i, bindGroup);
     }
 
-    if (x + y + z > 0) {
-        computePass.dispatchWorkgroups(x, y, z);
+    if (buffer !== undefined) {
+        computePass.dispatchWorkgroupsIndirect(buffer.underlyingResource(), <number>offset);
+    } else {
+        if (<number>x + <number>y + <number>z > 0) {
+            computePass.dispatchWorkgroups(<number>x, <number>y, <number>z);
+        }
     }
     computePass.end();