Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add compute dispatch indirect API #14970

Merged
merged 4 commits into from
Apr 16, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
24 changes: 21 additions & 3 deletions packages/dev/core/src/Compute/computeShader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import type { BaseTexture } from "../Materials/Textures/baseTexture";
import { Texture } from "../Materials/Textures/texture";
import { UniqueIdGenerator } from "../Misc/uniqueIdGenerator";
import type { IComputeContext } from "./IComputeContext";
import type { StorageBuffer } from "../Buffers/storageBuffer";
import { StorageBuffer } from "../Buffers/storageBuffer";
import { Logger } from "../Misc/logger";
import { TextureSampler } from "../Materials/Textures/textureSampler";
import type { DataBuffer } from "core/Buffers/dataBuffer";
Expand Down Expand Up @@ -348,6 +348,20 @@ 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 {
return this._dispatchInternal({ x, y, z });
}

/**
* 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 {
return this._dispatchInternal({ buffer, offset });
}

private _dispatchInternal(options: { x: number; y?: number; z?: number } | { buffer: StorageBuffer | DataBuffer; offset: number }): boolean {
stefnotch marked this conversation as resolved.
Show resolved Hide resolved
if (!this.fastMode) {
if (!this.isReady()) {
return false;
Expand Down Expand Up @@ -402,8 +416,12 @@ export class ComputeShader {
}
}

this._engine.computeDispatch(this._effect, this._context, this._bindings, x, y, z, this._options.bindingsMapping, this.gpuTimeInFrame);

if ("buffer" in options) {
Popov72 marked this conversation as resolved.
Show resolved Hide resolved
const buffer = options.buffer instanceof StorageBuffer ? options.buffer.getBuffer() : options.buffer;
stefnotch marked this conversation as resolved.
Show resolved Hide resolved
this._engine.computeDispatchIndirect(this._effect, this._context, this._bindings, buffer, offset, this._options.bindingsMapping, this.gpuTimeInFrame);
} else {
this._engine.computeDispatch(this._effect, this._context, this._bindings, options.x, options.y, options.z, this._options.bindingsMapping, this.gpuTimeInFrame);
}
return true;
}

Expand Down
32 changes: 32 additions & 0 deletions packages/dev/core/src/Engines/Extensions/engine.computeShader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import type { IComputePipelineContext } from "../../Compute/IComputePipelineCont
import { ThinEngine } from "../../Engines/thinEngine";
import type { Nullable } from "../../types";
import type { WebGPUPerfCounter } from "../WebGPU/webgpuPerfCounter";
import type { DataBuffer } from "../../Buffers/dataBuffer";

/**
* Type used to locate a resource in a compute shader.
Expand Down Expand Up @@ -113,6 +114,27 @@ declare module "../../Engines/thinEngine" {
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
Expand Down Expand Up @@ -170,6 +192,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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,21 @@ 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 */
_computeDispatch(
effect: ComputeEffect,
context: IComputeContext,
bindings: ComputeBindingList,
options: { x: number; y: number; z: number } | { buffer: DataBuffer; offset: number },
stefnotch marked this conversation as resolved.
Show resolved Hide resolved
bindingsMapping?: ComputeBindingMapping,
gpuPerfCounter?: WebGPUPerfCounter
): void;
}
}

Expand Down Expand Up @@ -67,6 +77,29 @@ WebGPUEngine.prototype.computeDispatch = function (
z = 1,
bindingsMapping?: ComputeBindingMapping,
gpuPerfCounter?: WebGPUPerfCounter
): void {
this._computeDispatch(effect, context, bindings, { x, y, z }, 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, { buffer, offset }, bindingsMapping, gpuPerfCounter);
};

WebGPUEngine.prototype._computeDispatch = function (
stefnotch marked this conversation as resolved.
Show resolved Hide resolved
effect: ComputeEffect,
context: IComputeContext,
bindings: ComputeBindingList,
options: { x: number; y: number; z: number } | { buffer: DataBuffer; offset: number },
bindingsMapping?: ComputeBindingMapping,
gpuPerfCounter?: WebGPUPerfCounter
): void {
this._endCurrentRenderPass();

Expand Down Expand Up @@ -97,8 +130,12 @@ WebGPUEngine.prototype.computeDispatch = function (
computePass.setBindGroup(i, bindGroup);
}

if (x + y + z > 0) {
computePass.dispatchWorkgroups(x, y, z);
if ("buffer" in options) {
computePass.dispatchWorkgroupsIndirect(options.buffer.underlyingResource(), options.offset ?? 0);
} else {
if (options.x + options.y + options.z > 0) {
computePass.dispatchWorkgroups(options.x, options.y, options.z);
}
}
computePass.end();

Expand Down