Playing with webgpu in three.js

2024-10-31

These are some old demos of basic three.js. See demos in three

I want to start learning three.js vertex and fragment shaders. Three.js has very early support for WebGPU, which has its own shader language called WGSL. Three also created their own shader language called TSL, which compiles to WGSL. There isn’t much documentation on either, which is why I am writing this post.

I start by adapting a boilerplate from Three.js Tutorials but i will be using the WebGPURenderer.

import * as THREE from 'three/webgpu';

const scene: THREE.Scene = new THREE.Scene();

// Use orthographic camera for 2D
const camera: THREE.OrthographicCamera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0.1, 10);
camera.position.z = 1;

// Define fixed width and height for canvas
const container: HTMLDivElement | null = document.getElementById('canvas-container') as HTMLDivElement;
if (!container) {
    throw new Error('Canvas container not found');
}
const width: number = 512;
const height: number = 512;

// Create a renderer
const renderer: THREE.WebGPURenderer = new THREE.WebGPURenderer();
renderer.setSize(width, height);
container.appendChild(renderer.domElement);

// Render the scene
function animate(): void {
    // Any updates to the scene go here
    
    renderer.render(scene, camera);
}

renderer.setAnimationLoop(animate);

This is the code I used to create a red triangle in WebGPU. If you want to learn pure WebGPU, I recommend WebGPU Fundamentals. Basically, what we need is a vertex shader and a fragment shader. The vertex shader is responsible for transforming the vertices of the triangle, and the fragment shader is responsible for coloring the triangle.

// fragment shader
@fragment
fn main() -> @location(0) vec4f {
  return vec4(1.0, 0.0, 0.0, 1.0);
}

// vertex shader
@vertex
fn main(
  @builtin(vertex_index) VertexIndex : u32
) -> @builtin(position) vec4f {
  var pos = array<vec2f, 3>(
    vec2(0.0, 0.5),
    vec2(-0.5, -0.5),
    vec2(0.5, -0.5)
  );

  return vec4f(pos[VertexIndex], 0.0, 1.0);
}

We can use the WGSL code directly in three by using the wgslFn function. I did have to make some changes to the code to make it work with three.js.

// Step 1: Define a simple material
const material = new THREE.MeshBasicNodeMaterial();

// Step 2: Define a fragment shader that returns a red color
const fragmentShader = wgslFn(`
fn frag() -> vec4f {
    return vec4(1.0, 0.0, 0.0, 1.0); // Red color with full opacity
}
`);
material.fragmentNode = fragmentShader();

// Step 3: Define a vertex shader
const vUv = varyingProperty("vec2", "vUv"); // Define a varying property for UV coordinates
const vertexShader = wgslFn(
    `
    fn crtVertex(
        position: vec3f,
        uv: vec2f
    ) -> vec3<f32> {
        varyings.vUv = uv; // Pass UV coordinates to the fragment shader
        return position; // Return the vertex position
    }
    `,
    [vUv]
);
material.positionNode = vertexShader({
    position: attribute("position"), // Attribute for vertex positions
    uv: attribute("uv") // Attribute for UV coordinates
});

// Step 4: Create geometry for a triangle
const vertices = new Float32Array([
    -0.5, -0.5, 0, // Vertex 1
    0.5, -0.5, 0,  // Vertex 2
    0, 0.5, 0      // Vertex 3
]);
const geometry = new THREE.BufferGeometry();
geometry.setAttribute('position', new THREE.BufferAttribute(vertices, 3)); // Set position attribute

// Step 5: Create a mesh with the geometry and material, and add it to the scene
const triangle = new THREE.Mesh(geometry, material);
scene.add(triangle); // Add the triangle mesh to the scene

Instead of defining the vertex positions inside the vertex shader, I defined them in the geometry and passed them in using the position attribute. I also had to define a varying property for the UV coordinates and pass them in using the uv attribute.

It doesnt look like the vertex shader is doing anything. I could get rid of it and the result would be the same. But in the future it will be useful for more complex shaders. Lets do the same with TSL.

// Step 1: Define a simple material
const material = new THREE.MeshBasicNodeMaterial();

// Step 2: Define the color of the material
material.fragmentNode = color(vec4(0.0, 1.0, 0.0, 1.0)); // Green!

// Step 3: Define positionNode
const vertexShader = Fn(() => {
    vUv.assign(uv());
    return THREE.positionGeometry;
})
material.positionNode = vertexShader();

Seems much simpler. Ive put the two examples side by side above. Three.js has an editor that converts TSL to WGSL. I couldn’t find much information on converting WGSL to TSL, so i’ll link some examples I found below.