webgpu-app/compute.js

146 lines
4.1 KiB
JavaScript

import { getPath } from "./common.js";
class Compute {
constructor(device, canvasFormat, wgsl)
{
const label = "compute";
this.label = label;
//////////////////////////////////////////////////////////////////////
// shader module
//////////////////////////////////////////////////////////////////////
const shaderModule = device.createShaderModule({
label: `${label} shader`,
code: wgsl,
});
//////////////////////////////////////////////////////////////////////
// bind group layout
//////////////////////////////////////////////////////////////////////
this.bindGroupLayout = {
compute: device.createBindGroupLayout({
label: `${label} bind group layout`,
entries: [
{
binding: 0,
visibility: GPUShaderStage.COMPUTE,
storageTexture: {
access: "write-only",
format: canvasFormat,
viewDimension: "2d",
}
}
]
}),
};
//////////////////////////////////////////////////////////////////////
// pipeline
//////////////////////////////////////////////////////////////////////
const pipelineLayout = device.createPipelineLayout({
label: `${label} pipeline layout`,
bindGroupLayouts: [
this.bindGroupLayout.compute,
],
});
this.computePipeline = device.createComputePipeline({
label: `${label} compute pipeline`,
layout: pipelineLayout,
compute: {
module: shaderModule,
},
});
this.textures = undefined;
}
recreateTextures(device, canvasTexture)
{
const label = this.label;
if (this.textures !== undefined && this.textures[0].width == canvasTexture.width && this.textures[0].height == canvasTexture.height) {
return;
}
if (this.textures !== undefined) {
for (let i = 0; i < 2; i++) {
this.textures[i].destroy();
}
}
//////////////////////////////////////////////////////////////////////
// textures
//////////////////////////////////////////////////////////////////////
this.textures = [];
for (let i = 0; i < 2; i++) {
const texture = device.createTexture({
label: `${label} texture ${i}`,
format: canvasTexture.format,
usage: GPUTextureUsage.STORAGE_BINDING | GPUTextureUsage.COPY_SRC,
size: [canvasTexture.width, canvasTexture.height],
});
this.textures.push(texture);
}
//////////////////////////////////////////////////////////////////////
// bind group
//////////////////////////////////////////////////////////////////////
this.bindGroups = []
for (let i = 0; i < 2; i++) {
const bindGroup = device.createBindGroup({
label: `${label} bind group ${i}`,
layout: this.bindGroupLayout.compute,
entries: [{
binding: 0,
resource: this.textures[i].createView(),
}],
});
this.bindGroups.push(bindGroup);
}
}
render(device, frameNumber, canvasTexture)
{
this.recreateTextures(device, canvasTexture);
const label = this.label;
const encoder = device.createCommandEncoder({ label: `${label} command encoder` });
const computePass = encoder.beginComputePass({ label: `${label} compute pass` });
computePass.setPipeline(this.computePipeline);
computePass.setBindGroup(0, this.bindGroups[frameNumber]);
computePass.dispatchWorkgroups(canvasTexture.width / 16, canvasTexture.height / 16);
computePass.end();
encoder.copyTextureToTexture({
texture: this.textures[frameNumber],
}, {
texture: canvasTexture
}, {
width: canvasTexture.width,
height: canvasTexture.height,
depthOrArrayLayers: 1,
});
const commandBuffer = encoder.finish();
device.queue.submit([commandBuffer]);
}
}
async function loadCompute(device, canvasFormat)
{
const canvasStorage = `@group(0) @binding(0) var out: texture_storage_2d<${canvasFormat}, write>;\n`
const wgsl = canvasStorage + await getPath("compute.wgsl");
const compute = new Compute(device, canvasFormat, wgsl);
return compute;
}
export { loadCompute };