WebGPU compute shader support prototype.

[takahirox]

This draft PR adds initial Compute shader support. I'm not sure if you want to add now (because you might like to wait for node material or WGSL, or might think we should focus on basic rendering features) so I made this PR as draft so far.

Even if we go with this PR, I think there is big room for further simplifying, robustness, and flexibility, so feedback is very welcome.

Similar to #20319, I'm happy if this suggestion helps improve webgpu renderer design and/or we partially adopt some codes even if we decide not to support compute shader yet.

Live demo

https://raw.githack.com/takahirox/three.js/ComputeShader/examples/#webgpu_compute

I recommend to run in small window because point size is one pixel and hard to see.

API and example

I add WebGPURenderer.compute() which takes an array of object having

{
  num: used for compute pass .dispatch(),
  shader: compute shader code strings,
  bindings: an array of webgpu binding
}

And I add WebGPUStorageBuffer corresponding to buffer in shader which can be input/output of compute shader. It takes Three.js attribute and webgpu buffer is created for it in renderer. So, if you add the attribute to geometry you can use computed data for rendering.

Currently you need to write raw shader code and manually create bindings but I expect we use node based material system even for compute shader and node builder will build the shader and bindings if node based material system lands.

This is an example computing particle positions in compute shader and then rendering them as Points.

const particleNum = 1024;
const particleSize = 3;

const particleArray = new Float32Array( particleNum * particleSize );
const particleAttribute = new THREE.BufferAttribute( particleArray, particleSize );
const particleBuffer = new WebGPUStorageBuffer( 'particle', particleAttribute );

const computeShader = `#version 450
  #define PARTICLE_NUM ${particleNum}
  #define PARTICLE_SIZE ${particleSize}

  layout(set = 0, binding = 0) buffer Particle {
    float particle[ PARTICLE_NUM * PARTICLE_SIZE ];
  } particle;

  void main() {
    uint index = gl_GlobalInvocationID.x;
    // compute and update buffer
    particle.particle[ index * PARTICLE_SIZE + 0 ] += bar0;
    particle.particle[ index * PARTICLE_SIZE + 1 ] += bar1;
    particle.particle[ index * PARTICLE_SIZE + 2 ] += bar2;
  }
`;

const computeParams = [ {
  num: particleNum,
  shader: computeShader,
  bindings: [ particleBuffer ]
} ];

// Create Points with the particle attribute to share WebGPU buffer

const pointsGeometry = new BufferGeometry().setAttribute(
  'position', particleAttribute
);
const pointsMaterial = new PointsMaterial( {} );
const points = new Points( pointsGeometry, pointsMaterial );
scene.add( points );

function render() {
  // compute particles and then render them
  renderer.compute( computeParams );
  renderer.render( scene, camera );
}

See webgpu_compute.html for more detail.

Changes and notes

WebGPURenderer

• Adding .compute() to WebGPURenderer rather than making new class like WebGPUComputer because I want computed data to be abled to be used for rendering by sharing WebGPU resources.
• I haven't added yet but a method which loads gpu buffer data to attribute array would be necessary

WebGPUComputePipelines (new)

• It manages compute pipelines. I think adding WebGPUComputePipelines would be simpler than adding compute pipelines management to WebGPURenderPipelines. (And I guess it's expected so named as WebGPU"Render"Pipelines.).

WebGPUStorageBuffer (new)

• New binding corresponding to buffer in the shader.

WebGPUBindings

• Adding getForCompute(). It would be simpler than letting the existing get() handle bindings for compute.
• Adding WebGPUStorageBuffer support.
• It needs to create webgpu buffer for WebGPUStorageBuffer which can be reused for rendering so it has dependency with WebGPUAttributes now.

WebGPUAttributes

• update() takes the new third argument customUsage because webgpu buffer for outputting in compute shader needs GPUBufferUsage.STORAGE. It saves and checks usage of buffer. If new usage buffer for a certain attribute is requested it destroys the existing buffer and creates a new one. (I'm not sure if this design is good tho...)

webgpu_compute

• An example showing how to use compute shader. Please feel free to rewrite with the one looking cooler.

[Mugen87]

This looks very nice! I think you can mark it as "Ready for review". I vote to merge this so there is a showcase for compute shaders.

[takahirox]

Thanks. I switched to "ready for review".

BTW, can we switch point color from red to white? webgpu_compute would look nicer.

[Kangz]

Note that it might not be super useful to have multiple GPUBindGroups for compute. The reason why there are multiple bindgroups is that we want to support 10k+ draws per frame in WebGPU so you can reuse most bindgroups and change only a small part of the bindings per-objects. The assumption is less true for compute shaders and my guesstimate is that you'd rarely need to have more than 100 per frame so it's probably ok to create bindgroups dynamically. WDYT?

[Mugen87]

At the current state of the renderer, resources like pipelines and bind groups are not yet shared across objects. We've planned to refactor this part when the initial version of the new node-based material system is integrated.

The assumption is less true for compute shaders and my guesstimate is that you'd rarely need to have more than 100 per frame so it's probably ok to create bindgroups dynamically.

I'm not sure I understand this bit, sorry. Does the current implementation not already do this?

[Kangz]

I'm not sure I understand this bit, sorry. Does the current implementation not already do this?

Whoops, that wasn't very clear, basically you probably don't need more than one bindgroup per compute shader (since multiple bindgroups is used to optimize high-throughput cases), so WebGPU.compute could take an array of buffer / textures instead of an array of bindings. This could simplify the API slightly, maybe.

The API as it is now looks great too, so if this is confusing feel free to ignore my comment.

[Mugen87]

How about starting with this implementation for now and try your suggestion with an additional PR? I think it would be useful to investigate your approach in order to find out which one works better.

In general, myself and potentially other devs understand a suggestion easier when seeing the concrete code changes. Especially since compute shaders and WebGPU in general are new topics for me (and the project)^^.

[Kangz]

Indeed, I should have showed some code, and now that I tried to do it, I just realized my suggestion was non-sensical. Apologies for the noise.

[Mugen87]

No problem. Your support here is highly appreciated 😊.

[takahirox]

Thanks for the comment. The purpose of accepting multiple bindgroups is for rendering computed result. Currently webgpu renderer supports only single bindgroup per pipeline. bindings takes an array for single bindgrouop, so for example you can put multiple buffers, uniforms, textures, and samplers in a single bindgroup.

Currently webgpu renderer creates webgpu buffer per an Three.js attribute. The buffer is used from offset 0. So, for example if user wants to compute position, uv, and normal they need to be written to different buffers (bindgroups) respectively.

If webgpu renderer starts to support interleaved buffer attribute (or something new packed attribute(?)) it would be resolved and single bindgroup would be good enough.

And I expect we will use node based material system even for compute shader if it lands, and node builder builds shader and bindgroups then most of bindings API would be hidden from user.

I'm still new to webgpu and compute shader. Please let me know if I mentioned anything weird.

I hope we keep investigation and discussion especially when we refactor and optimize webgpu renderer with node based material system or interleaved buffer attribute.

[takahirox]

Some thoughts about this PR.

1. Compute shader updates webgpu buffer data then Three.js attribute.array data and webgpu buffer data won't be synched. Might it be confusing to users? Would writing a note to document be good enough?

2. Some users may like to use computed result in CPU(JS). We may need to provide a way to load webgpu buffer data to Three.js attribute array. Two options in my mind. Adding new usage to BufferAttribute.usage always synching with webgpu buffer data by loading to attribute array each frame. It would be easy for users and also resolves 1 but it would be costly. Another one is adding a new method to webgpu renderer like .loadBuffer( attribute ).

3. Please feel free to refactor and clean up the code and APIs. I'm still new to webgpu and compute shader so maybe they can be better.

4. Please feel free to rewrite webgpu_compute example with the one looking cooler because I'm not an artist. And changing point color from red to white in WebGPURenderPipeline would make it nicer.

[Kangz]

Currently webgpu renderer creates webgpu buffer per an Three.js attribute. The buffer is used from offset 0. So, for example if user wants to compute position, uv, and normal they need to be written to different buffers (bindgroups) respectively.

This seems to hint that you think a single buffer is allowed per bindgroup. The goal of bindgroups is to batch multiple resource updates together so actually you are encouraged to put multiple buffers (and textures and samplers) per bindgroup.

And I expect we will use node based material system even for compute shader if it lands, and node builder builds shader and bindgroups then most of bindings API would be hidden from user.

Makes sense to hide such implementation details from the user.

[takahirox]

Sorry, I noticed that seems like I had mistaken bindgroup for bindgroup entry in #20390 (comment).

Currently webgpu renderer supports only single bind group per a pipeline regardless rendering or computing, and it can put multiple buffers per single bindgroup. webgpu_compute example actually does.