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

Open-source WebGL/WebGPU 3D engine for the web

import { BindGroup } from '../bind-group.js';
import { UniformBuffer } from '../uniform-buffer.js';

const _indirectDispatchEntryByteSize = 3 * 4;
class WebgpuCompute {
		destroy() {
				this.uniformBuffers.forEach((ub)=>ub.destroy());
				this.uniformBuffers.length = 0;
				this.bindGroup.destroy();
				this.bindGroup = null;
		}
		updateBindGroup() {
				const { bindGroup } = this;
				bindGroup.updateUniformBuffers();
				bindGroup.update();
		}
		dispatch(x, y, z) {
				const device = this.compute.device;
				device.setBindGroup(0, this.bindGroup);
				const passEncoder = device.passEncoder;
				passEncoder.setPipeline(this.pipeline);
				const { indirectSlotIndex, indirectBuffer, indirectFrameStamp } = this.compute;
				if (indirectSlotIndex >= 0) {
						let gpuBuffer;
						if (indirectBuffer) {
								gpuBuffer = indirectBuffer.impl.buffer;
						} else {
								gpuBuffer = device.indirectDispatchBuffer.impl.buffer;
						}
						const offset = indirectSlotIndex * _indirectDispatchEntryByteSize;
						passEncoder.dispatchWorkgroupsIndirect(gpuBuffer, offset);
				} else {
						passEncoder.dispatchWorkgroups(x, y, z);
				}
		}
		constructor(compute){
				this.uniformBuffers = [];
				this.bindGroup = null;
				this.compute = compute;
				const { device, shader } = compute;
				const { computeBindGroupFormat, computeUniformBufferFormats } = shader.impl;
				this.bindGroup = new BindGroup(device, computeBindGroupFormat);
				if (computeUniformBufferFormats) {
						for(const name in computeUniformBufferFormats){
								if (computeUniformBufferFormats.hasOwnProperty(name)) {
										const ub = new UniformBuffer(device, computeUniformBufferFormats[name], true);
										this.uniformBuffers.push(ub);
										this.bindGroup.setUniformBuffer(name, ub);
								}
						}
				}
				this.pipeline = device.computePipeline.get(shader, computeBindGroupFormat);
		}
}

export { WebgpuCompute };