TSL Compute Shaders: Introduction

1. Basic Example: Your First Compute Shader

Let's start with a simple compute shader that multiplies an array of numbers by 2. This demonstrates the fundamental workflow of TSL compute shaders. We'll build it step-by-step.

Step 1: Setup Renderer and Buffers

The first step in any GPU computation is to prepare the data. We need a WebGPU renderer to communicate with the GPU and buffers to hold our data.

// Initialize WebGPU renderer
const renderer = new THREE.WebGPURenderer()
await renderer.init()

// Create buffers for 10 float values
const count = 10
const inputBuffer = instancedArray(count, 'float')
const outputBuffer = instancedArray(count, 'float')

Step 2: Define the Compute Logic

With our buffers ready, we define the actual computation to be performed on the GPU. This is done by creating a TSL function.

// Main computation: multiply each value by 2
const multiplyCompute = Fn(() => {
  const input = inputBuffer.element(instanceIndex)
  const output = outputBuffer.element(instanceIndex)
  output.assign(input.mul(2))
})()

We also need a small function to initialize our input data with values from 1 to 10. This uses the same principles.

// Initialize input data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
const initCompute = Fn(() => {
  const input = inputBuffer.element(instanceIndex)
  input.assign(instanceIndex.add(1).toFloat())
})()

Step 3: Execute and Retrieve Data

Now we tell the GPU to run our compiled functions and then we retrieve the results.

// Execute compute shaders
await renderer.computeAsync(initCompute.compute(count))
await renderer.computeAsync(multiplyCompute.compute(count))

// Read results back to CPU
const inputArray = await renderer.getArrayBufferAsync(inputBuffer.value)
const outputArray = await renderer.getArrayBufferAsync(outputBuffer.value)

Here is the live output from the running code. As you can see, each number in the input array has been successfully multiplied by 2.

And here is the complete, self-contained function that accomplishes the task. The code below is what's actually running on this page.

import * as THREE from 'three/webgpu'
import { Fn, instancedArray, instanceIndex } from 'three/tsl'

async function initComputeShader() {
  // Initialize WebGPU renderer
  const renderer = new THREE.WebGPURenderer()
  await renderer.init()
  
  // Create buffers for 10 float values
  const count = 10
  const inputBuffer = instancedArray(count, 'float')
  const outputBuffer = instancedArray(count, 'float')
  
  // Initialize input data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
  const initCompute = Fn(() => {
    const input = inputBuffer.element(instanceIndex)
    input.assign(instanceIndex.add(1).toFloat())
  })()
  
  // Main computation: multiply each value by 2
  const multiplyCompute = Fn(() => {
    const input = inputBuffer.element(instanceIndex)
    const output = outputBuffer.element(instanceIndex)
    output.assign(input.mul(2))
  })()
  
  // Execute compute shaders
  await renderer.computeAsync(initCompute.compute(count))
  await renderer.computeAsync(multiplyCompute.compute(count))
  
  // Read results back to CPU
  const inputArray = await renderer.getArrayBufferAsync(inputBuffer.value)
  const outputArray = await renderer.getArrayBufferAsync(outputBuffer.value)
  
  return {
    input: Array.from(new Float32Array(inputArray)),
    output: Array.from(new Float32Array(outputArray))
  }
}

2. Game of Life: 2D Grid Simulation

Now let's see something more complex: a complete implementation of Conway's Game of Life running entirely on the GPU. This showcases 2D array processing, neighbor counting, and conditional logic, all built step-by-step.

Step 1: Setup Grid and Buffers

We start by setting up our simulation environment. This includes defining the grid size and creating two GPU buffers with instancedArray: one to hold the current state of the cells (currentGeneration) and another for the next state (nextGeneration). We use two buffers to avoid race conditions, where a cell's new state might incorrectly influence its neighbors' calculations in the same step.

// Grid dimensions - 64x64 = 4,096 cells total
const gridWidth = 64;
const gridHeight = 64;
const totalCells = gridWidth * gridHeight;

// Create buffers for current and next generation (using integers: 0 for dead, 1 for alive)
const currentGeneration = instancedArray(totalCells, 'int');
const nextGeneration = instancedArray(totalCells, 'int');

Step 2: 2D to 1D Mapping

GPU buffers are linear, one-dimensional arrays. To simulate a 2D grid, we need helper functions to convert (x, y) coordinates into a 1D index. We also create a helper to get a cell's state with "toroidal" or "wrapping" boundaries, where the grid's edges connect to each other.

// Helper function to convert 2D coordinates to a 1D index
const getIndex = Fn(([x, y]) => {
  return y.mul(gridWidth).add(x);
});

// Helper function to get cell state with boundary wrapping
const getCell = Fn(([buffer, x, y]) => {
  // Wrap coordinates for toroidal topology (edges connect)
  const wrappedX = x.add(gridWidth).mod(gridWidth);
  const wrappedY = y.add(gridHeight).mod(gridHeight);
  const index = getIndex(wrappedX, wrappedY);
  return buffer.element(index);
});

Step 3: Initializing the Grid

To start the simulation, we need an initial pattern. We'll write a compute shader that gives each cell a 30% chance of being "alive." We use TSL's built-in hash() function to generate a pseudo-random value for each cell based on its instanceIndex.

// Initialize grid with a random pattern
const initializeGrid = Fn(() => {
  const currentCell = currentGeneration.element(instanceIndex);

  // Use hash function for pseudo-random initialization
  const randomValue = hash(instanceIndex.add(12345)); // Add a seed

  // 30% chance for a cell to be alive initially
  If(randomValue.lessThan(0.3), () => {
    currentCell.assign(1); // Alive
  }).Else(() => {
    currentCell.assign(0); // Dead
  });
})();

Step 4: The Game of Life Update Logic

This is the core of the simulation. For each cell, we count its eight living neighbors. Then, we apply the classic rules of Conway's Game of Life using TSL's If().ElseIf().Else() structure to determine if the cell should be alive or dead in the next generation.

const updateGeneration = Fn(() => {
  // Convert 1D thread index to 2D grid coordinates
  const x = instanceIndex.mod(gridWidth);
  const y = instanceIndex.div(gridWidth).toInt();

  // Count living neighbors
  const neighbors = int(0).toVar();
  neighbors.addAssign(getCell(currentGeneration, x.sub(1), y.sub(1)));
  neighbors.addAssign(getCell(currentGeneration, x.sub(1), y));
  neighbors.addAssign(getCell(currentGeneration, x.sub(1), y.add(1)));
  neighbors.addAssign(getCell(currentGeneration, x, y.sub(1)));
  neighbors.addAssign(getCell(currentGeneration, x, y.add(1)));
  neighbors.addAssign(getCell(currentGeneration, x.add(1), y.sub(1)));
  neighbors.addAssign(getCell(currentGeneration, x.add(1), y));
  neighbors.addAssign(getCell(currentGeneration, x.add(1), y.add(1)));

  const currentCell = currentGeneration.element(instanceIndex);
  const nextCell = nextGeneration.element(instanceIndex);

  // Apply Conway's Game of Life rules
  If(currentCell.equal(1), () => { // Current cell is alive
    If(neighbors.lessThan(2).or(neighbors.greaterThan(3)), () => {
      nextCell.assign(0); // Dies from under/overpopulation
    }).Else(() => {
      nextCell.assign(1); // Survives
    });
  }).Else(() => { // Current cell is dead
    If(neighbors.equal(3), () => {
      nextCell.assign(1); // Birth
    }).Else(() => {
      nextCell.assign(0); // Stays dead
    });
  });
})();

Step 5: Running the Simulation

Finally, we orchestrate the simulation. We first run the initializeGrid shader once to set up the initial state. Then, we provide a button to run the simulation step-by-step. Each click will execute the updateGeneration shader to calculate the new state, followed by a copyGeneration shader to copy the results from nextGeneration back to currentGeneration, preparing it for the next iteration.

// Copy next generation to current generation for the next iteration
const copyGeneration = Fn(() => {
  currentGeneration.element(instanceIndex).assign(nextGeneration.element(instanceIndex));
})();

// Initialize the grid with a random pattern
await renderer.computeAsync(initializeGrid.compute(totalCells));

// Run simulation for multiple generations
for (let step = 0; step < 10; step++) {
  // Calculate next generation and write to the second buffer
  await renderer.computeAsync(updateGeneration.compute(totalCells));
  // Copy the new state back to the first buffer for the next read
  await renderer.computeAsync(copyGeneration.compute(totalCells));
}

Live Simulation

Here is the live simulation. Press the button to advance the Game of Life by one generation and see the simulation evolve.

Complete Code

And here is the complete, self-contained function. The code below is what's actually running on this page.

import * as THREE from 'three/webgpu'
import { Fn, instancedArray, instanceIndex, int, hash, If } from 'three/tsl'

async function initGameOfLife() {
  // Initialize WebGPU renderer
  const renderer = new THREE.WebGPURenderer()
  await renderer.init()
  
  // Grid dimensions - 64x64 = 4,096 cells total
  const gridWidth = 64
  const gridHeight = 64
  const totalCells = gridWidth * gridHeight
  
  // Create buffers for current and next generation
  const currentGeneration = instancedArray(totalCells, 'int')
  const nextGeneration = instancedArray(totalCells, 'int')
  
  // Helper function to convert 2D coordinates to 1D index
  const getIndex = Fn(([x, y]: any) => {
    return y.mul(gridWidth).add(x)
  })
  
  // Helper function to get cell state with boundary wrapping
  const getCell = Fn(([buffer, x, y]: any) => {
    // Wrap coordinates for toroidal topology (edges connect)
    const wrappedX = x.add(gridWidth).mod(gridWidth)
    const wrappedY = y.add(gridHeight).mod(gridHeight)
    const index = getIndex(wrappedX, wrappedY)
    return buffer.element(index)
  })
  
  // Initialize grid with random pattern
  const initializeGrid = Fn(() => {
    const currentCell = currentGeneration.element(instanceIndex)
    
    // Use hash function for pseudo-random initialization
    const randomValue = hash(instanceIndex.add(12345))
    
    // 30% chance for a cell to be alive initially
    If(randomValue.lessThan(0.3), () => {
      currentCell.assign(1) // Alive
    }).Else(() => {
      currentCell.assign(0) // Dead
    })
  })()
  
  // Game of Life update logic - the heart of the simulation
  const updateGeneration = Fn(() => {
    // Convert 1D thread index to 2D grid coordinates
    const x = instanceIndex.mod(gridWidth)
    const y = instanceIndex.div(gridWidth).toInt()
    
    // Count living neighbors (all 8 surrounding cells)
    const neighbors = int(0).toVar()
    
    // Check all 8 neighboring cells manually
    // (TSL doesn't support dynamic loops over arrays)
    neighbors.addAssign(getCell(currentGeneration, x.sub(1), y.sub(1)))
    neighbors.addAssign(getCell(currentGeneration, x.sub(1), y))
    neighbors.addAssign(getCell(currentGeneration, x.sub(1), y.add(1)))
    neighbors.addAssign(getCell(currentGeneration, x, y.sub(1)))
    neighbors.addAssign(getCell(currentGeneration, x, y.add(1)))
    neighbors.addAssign(getCell(currentGeneration, x.add(1), y.sub(1)))
    neighbors.addAssign(getCell(currentGeneration, x.add(1), y))
    neighbors.addAssign(getCell(currentGeneration, x.add(1), y.add(1)))
    
    // Get current cell state and prepare next state
    const currentCell = currentGeneration.element(instanceIndex)
    const nextCell = nextGeneration.element(instanceIndex)
    
    // Apply Conway's Game of Life rules
    If(currentCell.equal(1), () => {
      // Current cell is alive
      If(neighbors.lessThan(2), () => {
        nextCell.assign(0) // Dies from underpopulation
      }).ElseIf(neighbors.greaterThan(3), () => {
        nextCell.assign(0) // Dies from overpopulation
      }).Else(() => {
        nextCell.assign(1) // Survives (2 or 3 neighbors)
      })
    }).Else(() => {
      // Current cell is dead
      If(neighbors.equal(3), () => {
        nextCell.assign(1) // Birth (exactly 3 neighbors)
      }).Else(() => {
        nextCell.assign(0); // Stays dead
      })
    })
  })()
  
  // Copy next generation to current generation for next iteration
  const copyGeneration = Fn(() => {
    const current = currentGeneration.element(instanceIndex)
    const next = nextGeneration.element(instanceIndex)
    current.assign(next)
  })()
  
  // Initialize the grid with random pattern
  await renderer.computeAsync(initializeGrid.compute(totalCells))
  
  // Run simulation for multiple generations
  for (let step = 0; step < 10; step++) {
    // Calculate next generation
    await renderer.computeAsync(updateGeneration.compute(totalCells))
    // Copy next generation to current
    await renderer.computeAsync(copyGeneration.compute(totalCells))
  }
  
  return {
    renderer,
    currentGeneration,
    nextGeneration,
    gridWidth,
    gridHeight,
    updateGeneration,
    copyGeneration
  }
}

We've run the entire Game of Life simulation on the GPU. Now, how do we see it? A key advantage of TSL is its seamless integration with the Three.js rendering pipeline. We can visualize our simulation results without ever needing to bring the data back to the CPU, allowing for high-performance, real-time graphics.

Step 1: From Buffers to Pixels with InstancedMesh

Our simulation involves 4,096 cells, and we need to draw a quad for each one. Creating and managing 4,096 separate objects would be very inefficient for the CPU. The solution is THREE.InstancedMesh, a special object that allows us to draw thousands of identical geometries in a single command, each with unique properties like position and color.

// Create scene and camera for a 2D orthographic view
const scene = new THREE.Scene();
const camera = new THREE.OrthographicCamera(-0.5, 0.5, 0.5, -0.5, 0.1, 10);
camera.position.z = 1;

// A single, small plane geometry will be used for all cells
const geometry = new THREE.PlaneGeometry(1 / gridWidth, 1 / gridHeight);

// Use an instanced mesh to draw all 4,096 cells in one efficient command
const mesh = new THREE.InstancedMesh(geometry, undefined, totalCells);
scene.add(mesh);

Step 2: Positioning Each Cell with a TSL Vertex Shader

Now that we have an InstancedMesh, we need to tell the GPU where to place each of the 4,096 plane instances. We do this with a TSL function assigned to the material's positionNode. This function is effectively a vertex shader that runs on the GPU. It uses the built-in instanceIndex to calculate a unique position for each instance, arranging them in a grid.

// A TSL-powered material
const material = new THREE.MeshBasicNodeMaterial();
mesh.material = material;

// The positionNode is a TSL function that runs for every vertex of every instance.
// It's a vertex shader written in TSL!
material.positionNode = Fn(() => {
  // Calculate 2D grid position (x, y) from the 1D instanceIndex
  const x = instanceIndex.mod(gridWidth);
  const y = instanceIndex.div(gridWidth).toInt();

  // Normalize coordinates to the range [-0.5, 0.5] to fit our camera view
  const uvX = x.toFloat().add(0.5).div(gridWidth).sub(0.5);
  const uvY = y.toFloat().add(0.5).div(gridHeight).sub(0.5);

  // Add this calculated per-instance offset to the geometry's local vertex position
  const finalPosition = positionLocal.add(vec4(uvX, uvY, 0, 0));

  return finalPosition;
})();

Step 3: Coloring Each Cell with a Compute Shader

This is the crucial step that connects our compute simulation to our visual rendering. We create a new compute shader whose sole job is to read the state of each cell from our currentGeneration buffer and write a corresponding color into a new colorBuffer. This colorBuffer is then fed directly into the material's colorNode, telling the rendering pipeline what color to make each instance.

// Step 3.1: Create a new buffer to hold per-instance color data
const colorBuffer = instancedArray(totalCells, 'vec4');

// Step 3.2: Create a compute shader to populate the color buffer
const updateColors = Fn(() => {
  const cellState = currentGeneration.element(instanceIndex);
  const outputColor = colorBuffer.element(instanceIndex);

  const aliveColor = vec4(0.0, 1.0, 0.0, 1.0); // Green
  const deadColor = vec4(0.0, 0.0, 0.0, 1.0);  // Black

  // If cell is alive (state == 1), output green. Otherwise, black.
  outputColor.assign(deadColor);
  If(cellState.equal(1), () => {
    outputColor.assign(aliveColor);
  });
})();

// Step 3.3: Pipe the color buffer directly into the material's color node
material.colorNode = colorBuffer.toAttribute();

Step 4: The Animation Loop: Tying It All Together

The final piece is the animation loop, which orchestrates all our GPU tasks frame by frame. On each frame, we execute our compute shaders in sequence and then render the final scene. All of this happens on the GPU, coordinated by a few asynchronous commands from the CPU.

// The main animation loop, managed by the renderer
renderer.setAnimationLoop(async () => {
  // 1. Run one step of the Game of Life simulation
  await runGameOfLifeStep(); // This runs updateGeneration and copyGeneration

  // 2. Update the color buffer based on the new simulation state
  await renderer.computeAsync(updateColors.compute(totalCells));

  // 3. Render the scene. The GPU now has all the data it needs.
  renderer.render(scene, camera);
});

With this, we have a complete, high-performance, GPU-powered simulation and visualization of Conway's Game of Life. The interactive canvas you saw in the previous section is rendered using exactly these techniques.

3. Langton's Ant: Emergent Complexity

Langton's Ant is a fascinating cellular automaton that demonstrates how simple rules can lead to complex emergent behavior. Despite having only two simple rules, the ant creates intricate patterns that transition from apparent randomness to structured "highways" after thousands of steps.

The Rules

The simulation follows just two simple rules:

• On a white cell: Turn right, flip the cell to black, then move forward
• On a black cell: Turn left, flip the cell to white, then move forward

Implementation

Our implementation uses a 200×200 grid (40,000 cells) to provide enough space for the ant to develop complex patterns. The simulation runs entirely on the GPU using TSL compute shaders for maximum performance.

4. Boids: Agent-Based Simulation

Next, we'll explore a more complex, dynamic simulation: Boids. Developed by Craig Reynolds, this is an artificial life program that simulates the flocking behavior of birds. Each "boid" follows a set of simple rules, and the combination of these rules leads to complex, emergent flocking behavior. It's a classic example of how simple local interactions can create complex global patterns, making it a perfect candidate for GPU acceleration.

The emergent flocking behavior is governed by three simple rules that each boid follows based on its perception of its local neighbors. These rules, illustrated below, are applied concurrently to every boid in the simulation.

The simulation uses three simple rules to generate complex flocking behavior:

• Separation: Steer to avoid crowding local flockmates. This prevents boids from clumping together and colliding.
• Alignment: Steer towards the average heading (direction of travel) of local flockmates. This helps the flock move as a cohesive group.
• Cohesion: Steer to move toward the average position (center of mass) of local flockmates. This keeps the flock together.

The fascinating aspect of Boids is that these simple, local rules, when applied to hundreds or thousands of agents, produce complex and life-like global flocking behavior. None of the boids have a concept of the entire flock, only their immediate surroundings. This makes the algorithm highly parallelizable and a perfect fit for a GPU compute shader, where we can calculate the behavior for every boid simultaneously.

Step 1: Configuration and State Storage

Like before, we start by defining the state of our simulation. For boids, each agent has a position and a velocity in 3D space. We also have several parameters that control their behavior.

// boids.ts
// Configuration for the simulation
export interface BoidsConfig {
  count: number; // Number of boids
  speedLimit: number;
  bounds: number;
  separation: number; // Radius for separation rule
  alignment: number;  // Radius for alignment rule
  cohesion: number;   // Radius for cohesion rule
  freedom: number;
}

With the parameters defined, we create buffers on the GPU to store the state for each boid:

// ... inside BoidsSimulation class ...

// Buffers to store boid state on the GPU
const positionArray = new Float32Array(count * 3);
const velocityArray = new Float32Array(count * 3);
const phaseArray = new Float32Array(count); // For animation

// ... initialize arrays with random data ...

// Create GPU buffers
const positionStorage = attributeArray(positionArray, 'vec3');
const velocityStorage = attributeArray(velocityArray, 'vec3');
const phaseStorage = attributeArray(phaseArray, 'float');

Step 2: The Rules of Flocking in a Compute Shader

The magic of boids comes from a set of simple steering behaviors that each boid follows based on its local neighborhood. We implement this entire logic in a single, large compute shader: computeVelocity. This shader calculates a new velocity for each boid based on several factors: the boid's desire to stay with the flock, avoid collisions, stay within the boundaries, and flee from predators. Let's break down how it works, part by part.

// Combine all rule distances into a single zone of influence
const zoneRadius = separation.add(alignment).add(cohesion).toConst();
// Calculate squared distance for cheaper comparisons
const zoneRadiusSq = zoneRadius.mul(zoneRadius).toConst();

// Pre-calculate normalized thresholds for switching between behaviors
const separationThresh = separation.div(zoneRadius).toConst();
const alignmentThresh = (separation.add(alignment)).div(zoneRadius).toConst();

// A vector pointing from the boid towards the world origin
const dirToCenter = position.toVar();
// Encourage horizontal flocking by weighting the y-axis
dirToCenter.y.mulAssign(2.5); 
// Apply a gentle force towards the center
velocity.subAssign(normalize(dirToCenter).mul(deltaTime).mul(5.0));

// Calculate the boid's squared distance to the ray
const directionToRay = rayOrigin.sub(position).toConst();
const projectionLength = dot(directionToRay, rayDirection).toConst();
const closestPoint = rayOrigin.sub(rayDirection.mul(projectionLength)).toConst();
const directionToClosestPoint = closestPoint.sub(position).toConst();
const distanceToClosestPoint = length(directionToClosestPoint).toConst();
const distanceToClosestPointSq = distanceToClosestPoint.mul(distanceToClosestPoint).toConst();

// If within the ray's radius, apply a repulsive force
const rayRadiusSq = float(150.0).mul(150.0).toConst();
If(distanceToClosestPointSq.lessThan(rayRadiusSq), () => {
  const velocityAdjust = (distanceToClosestPointSq.div(rayRadiusSq).sub(1.0)).mul(deltaTime).mul(100.0);
  velocity.addAssign(normalize(directionToClosestPoint).mul(velocityAdjust));
  limit.addAssign(5.0); // Temporarily increase speed limit to escape
});

// Loop through all other boids
Loop({ start: uint(0), end: uint(count), type: 'uint', condition: '<' }, ({ i }) => {
  // ... (skip self, get neighbor data) ...
  
  const dirToBird = birdPosition.sub(position);
  const distToBird = length(dirToBird);
  // ... (skip if distance is zero) ...
  const distToBirdSq = distToBird.mul(distToBird);

  // Is the neighbor within our zone of influence?
  If(distToBirdSq.greaterThan(zoneRadiusSq), () => {
    Continue(); // Skip boid, it's too far away
  });

  // 'percent' determines which rule to apply
  const percent = distToBirdSq.div(zoneRadiusSq);

  // ... apply rules based on 'percent' ...
});

// RULE 1: SEPARATION
If(percent.lessThan(separationThresh), () => {
  const velocityAdjust = (separationThresh.div(percent).sub(1.0)).mul(deltaTime);
  velocity.subAssign(normalize(dirToBird).mul(velocityAdjust));
});

// RULE 2: ALIGNMENT
}).ElseIf(percent.lessThan(alignmentThresh), () => {
  // 1. Find how far into the alignment zone the neighbor is (0.0 to 1.0)
  const threshDelta = alignmentThresh.sub(separationThresh);
  const adjustedPercent = (percent.sub(separationThresh)).div(threshDelta);

  // 2. Get the neighbor's velocity
  const birdVelocity = velocityStorage.element(i);

  // 3. Calculate a smooth weight based on the position in the zone.
  const cosRange = cos(adjustedPercent.mul(PI_2));
  const cosRangeAdjust = float(1.0).sub(cosRange.mul(0.5)); // Full formula: 1.0 - cos(percent * 2PI) * 0.5
  
  // 4. Apply the alignment force, scaled by the weight
  const velocityAdjust = cosRangeAdjust.mul(deltaTime);
  velocity.addAssign(normalize(birdVelocity).mul(velocityAdjust));
});

// RULE 3: COHESION
}).Else(() => {
  // 1. Find how far into the cohesion zone the neighbor is (0.0 to 1.0)
  const threshDelta = alignmentThresh.oneMinus();
  const adjustedPercent = threshDelta.equal(0.0).select(1.0, (percent.sub(alignmentThresh)).div(threshDelta));

  // 2. Calculate the same smooth weight as used in alignment.
  // This makes the cohesive force strongest in the middle of the zone.
  const cosRange = cos(adjustedPercent.mul(PI_2));
  const cosRangeAdjust = float(1.0).sub(cosRange.mul(0.5));
  
  // 3. Apply the cohesion force, steering towards the neighbor's position.
  const velocityAdjust = cosRangeAdjust.mul(deltaTime);
  velocity.addAssign(normalize(dirToBird).mul(velocityAdjust));
});

If(length(velocity).greaterThan(limit), () => {
  velocity.assign(normalize(velocity).mul(limit));
});

velocityStorage.element(birdIndex).assign(velocity);

Step 3: Updating Position

After all the velocity adjustments have been calculated, a second, much simpler compute shader called computePosition runs. It performs a basic physics update: the new position is the old position plus the new velocity, scaled by time.

// boids.ts -> computePosition
const computePosition = Fn(() => {
  // Standard physics update: position += velocity * deltaTime
  positionStorage.element(instanceIndex).addAssign(
    velocityStorage.element(instanceIndex).mul(deltaTime)
  );

  // ... (phase update logic for animation) ...
})();

Step 4: Running the Simulation

Finally, we orchestrate the simulation. In our main application loop, we call the two compute shaders in sequence. Although we don't have a visualization yet, we can run a test to verify the simulation is working. The test runs the simulation for a few frames and measures the average distance the boids have moved to confirm they are not static.

// In the main application loop...
boids.update(deltaTime); // Update uniforms like time
boids.compute(renderer); // This executes computeVelocity then computePosition

Step 1: Core Configuration and Initialization

We start by defining the configuration for our simulation. This includes the number of boids, the simulation boundaries, and the behavioral parameters for each species, such as separation, alignment, and cohesion. In TSL, these parameters are passed to the compute shader as `uniforms`.

Interspecies Dynamics: A Game of Rock-Paper-Scissors

To make the simulation more visually interesting, this version of Boids introduces three distinct species of boids, each with slightly different flocking parameters. More importantly, they interact with each other in a "Rock-Paper-Scissors" dynamic:

• Species 0 hunts Species 1.
• Species 1 hunts Species 2.
• Species 2 hunts Species 0.

When a boid encounters a member of a species it preys upon, it will exhibit hunting behavior (a strong cohesion force towards the prey). Conversely, when it encounters a predator, it will exhibit fleeing behavior (a strong separation force away from the predator). This creates a constantly shifting dynamic where different groups will chase and flee from each other, adding another layer of emergent complexity to the simulation.

Complete Code

And here is the complete code for the BoidsSimulation class. This contains all the setup, compute shader definitions, and update logic discussed above.

import * as THREE from 'three/webgpu';
import { 
  uniform, 
  attributeArray, 
  float, 
  uint, 
  Fn, 
  If, 
  Loop, 
  Continue, 
  normalize, 
  instanceIndex, 
  length, 
  dot, 
  cos, 
  max, 
  property,
} from 'three/tsl';

export interface BoidsConfig {
  count: number;
  speedLimit: number;
  bounds: number;
  separation: number;
  alignment: number;
  cohesion: number;
  freedom: number;
}

export class BoidsSimulation {
  // ... constructor and initialization ...
  
  private initializeCompute(): void {
    const { count, speedLimit } = this.config;
    const { positionStorage, velocityStorage, phaseStorage } = this.storage;

    const computeVelocity = Fn(() => {
      const PI = float(3.141592653589793);
      const PI_2 = PI.mul(2.0);
      const limit = property('float', 'limit').assign(speedLimit);

      const { alignment, separation, cohesion, deltaTime, rayOrigin, rayDirection } = this.uniforms;

      const zoneRadius = separation.add(alignment).add(cohesion).toConst();
      const separationThresh = separation.div(zoneRadius).toConst();
      const alignmentThresh = (separation.add(alignment)).div(zoneRadius).toConst();
      const zoneRadiusSq = zoneRadius.mul(zoneRadius).toConst();

      const birdIndex = instanceIndex.toConst('birdIndex');
      const position = positionStorage.element(birdIndex).toVar();
      const velocity = velocityStorage.element(birdIndex).toVar();

      const directionToRay = rayOrigin.sub(position).toConst();
      const projectionLength = dot(directionToRay, rayDirection).toConst();
      const closestPoint = rayOrigin.sub(rayDirection.mul(projectionLength)).toConst();
      const directionToClosestPoint = closestPoint.sub(position).toConst();
      const distanceToClosestPoint = length(directionToClosestPoint).toConst();
      const distanceToClosestPointSq = distanceToClosestPoint.mul(distanceToClosestPoint).toConst();

      const rayRadiusSq = float(150.0).mul(150.0).toConst();
      If(distanceToClosestPointSq.lessThan(rayRadiusSq), () => {
        const velocityAdjust = (distanceToClosestPointSq.div(rayRadiusSq).sub(1.0)).mul(deltaTime).mul(100.0);
        velocity.addAssign(normalize(directionToClosestPoint).mul(velocityAdjust));
        limit.addAssign(5.0);
      });

      const dirToCenter = position.toVar();
      dirToCenter.y.mulAssign(2.5);
      velocity.subAssign(normalize(dirToCenter).mul(deltaTime).mul(5.0));

      Loop({ start: uint(0), end: uint(count), type: 'uint', condition: '<' }, ({ i }) => {
        
        If(i.equal(birdIndex), () => {
          Continue();
        });

        const birdPosition = positionStorage.element(i);
        const dirToBird = birdPosition.sub(position);
        const distToBird = length(dirToBird);

        If(distToBird.lessThan(0.0001), () => {
          Continue();
        });

        const distToBirdSq = distToBird.mul(distToBird);

        If(distToBirdSq.greaterThan(zoneRadiusSq), () => {
          Continue();
        });

        const percent = distToBirdSq.div(zoneRadiusSq);

        If(percent.lessThan(separationThresh), () => {
          const velocityAdjust = (separationThresh.div(percent).sub(1.0)).mul(deltaTime);
          velocity.subAssign(normalize(dirToBird).mul(velocityAdjust));
        }).ElseIf(percent.lessThan(alignmentThresh), () => {
          const threshDelta = alignmentThresh.sub(separationThresh);
          const adjustedPercent = (percent.sub(separationThresh)).div(threshDelta);
          const birdVelocity = velocityStorage.element(i);

          const cosRange = cos(adjustedPercent.mul(PI_2));
          const cosRangeAdjust = float(1.0).sub(cosRange.mul(0.5));
          const velocityAdjust = cosRangeAdjust.mul(deltaTime);
          velocity.addAssign(normalize(birdVelocity).mul(velocityAdjust));
        }).Else(() => {
          const threshDelta = alignmentThresh.oneMinus();
          const adjustedPercent = threshDelta.equal(0.0).select(1.0, (percent.sub(alignmentThresh)).div(threshDelta));

          const cosRange = cos(adjustedPercent.mul(PI_2));
          const adj1 = cosRange.mul(-0.5);
          const adj2 = adj1.add(0.5);
          const adj3 = float(0.5).sub(adj2);

          const velocityAdjust = adj3.mul(deltaTime);
          velocity.addAssign(normalize(dirToBird).mul(velocityAdjust));
        });
      });

      If(length(velocity).greaterThan(limit), () => {
        velocity.assign(normalize(velocity).mul(limit));
      });

      velocityStorage.element(birdIndex).assign(velocity);
    })().compute(count);

    const computePosition = Fn(() => {
      const { deltaTime } = this.uniforms;
      positionStorage.element(instanceIndex).addAssign(velocityStorage.element(instanceIndex).mul(deltaTime).mul(15.0));

      const velocity = velocityStorage.element(instanceIndex);
      const phase = phaseStorage.element(instanceIndex);

      const modValue = phase.add(deltaTime).add(length(velocity.xz).mul(deltaTime).mul(3.0)).add(max(velocity.y, 0.0).mul(deltaTime).mul(6.0));
      phaseStorage.element(instanceIndex).assign(modValue.mod(62.83));
    })().compute(count);
    
    // ...
  }
}

Step 5: Visualizing the Flock

The simulation logic is complete, but how do we see the result? Visualizing thousands of dynamic agents efficiently is a challenge perfectly suited for the GPU. This section breaks down how the boids are rendered using a seamless, high-performance pipeline that connects the compute simulation directly to the graphics display without ever needing to bring data back to the CPU.

In-Depth: The Vertex Shader (boidVertexShader)

This is where most of the magic happens. The vertex shader runs for every vertex of the base triangle, for every boid instance. Its goal is to apply the correct position and orientation to each boid. Let's break down its execution step-by-step for a single vertex of a single boid:

// boids-visualization.ts
const boidVertexShader = Fn(() => {
  // Get the current instance index (which boid we're rendering)
  const boidIndex = instanceIndex;
  
  // Get position and velocity from compute storage
  const boidPosition = this.storage.positionStorage.element(boidIndex);
  const boidVelocity = this.storage.velocityStorage.element(boidIndex);
  
  // Transform local vertex position
  const localPos = positionLocal.toVar();
  localPos.mulAssign(this.config.particleSize);
  
  // Create a rotation matrix to align the boid with its velocity
  const velocity = normalize(boidVelocity.add(vec3(0.001, 0.001, 0.001))); // Add epsilon to avoid zero velocity

  const forward = velocity.toVar('forward');
  const up = vec3(0.0, 1.0, 0.0).toVar('up');
  const right = normalize(cross(up, forward)).toVar('right');
  const newUp = normalize(cross(forward, right)).toVar('newUp');
  
  const rotationMatrix = mat3(
    right.x, forward.x, newUp.x,
    right.y, forward.y, newUp.y,
    right.z, forward.z, newUp.z
  ).toVar();
  
  const rotatedPos = rotationMatrix.mul(localPos);
  
  // Translate to boid position
  const worldPos = rotatedPos.add(boidPosition);
  
  // Transform to clip space
  return cameraProjectionMatrix.mul(cameraViewMatrix).mul(vec4(worldPos, 1.0));
});

1. Identify the Boid: The shader gets the special instanceIndex variable, a unique ID from 0 to count-1 that tells it which boid instance it's currently processing.
2. Fetch State from GPU Storage: It uses the boidIndex to read that specific boid's position and velocity directly from the positionStorage and velocityStorage buffers. For example, this.storage.positionStorage.element(boidIndex) is the TSL code that translates to a GPU memory read at the correct offset in the buffer.
3. Orient the Boid (Vector Math): To make the boid triangle "point" in the direction it's flying, the shader constructs a 3D rotation matrix on-the-fly.
   • The boid's normalized velocity vector serves as the "forward" direction.
   • It calculates the "right" direction by taking the cross product of a world "up" vector `(0,1,0)` and the `forward` vector. The cross product of two vectors yields a third vector that is perpendicular to both.
   • To ensure the basis is perfectly orthogonal, it recalculates a `newUp` vector by taking the cross product of the `forward` and `right` vectors.
   • These three vectors (`right`, `newUp`, and `forward`) form the columns of a `mat3` rotation matrix.
4. Apply Rotation and Position: The original vertex position of the base triangle (positionLocal) is first scaled, then multiplied by the rotationMatrix. This orients the triangle in 3D space. The result is then added to the boid's unique world position fetched from the storage buffer.
5. Project to Screen: The final worldPos is multiplied by the camera's view and projection matrices. This is the standard 3D graphics transformation that converts the 3D world coordinate into the 2D coordinate that will be displayed on the screen.

In-Depth: The Fragment Shader (boidFragmentShader)

After the vertex shader has positioned the triangle on the screen, the fragment (or pixel) shader runs for every pixel inside that triangle. Its job is to determine the pixel's color.

// boids-visualization.ts
const boidFragmentShader = Fn(() => {
  const boidIndex = instanceIndex;
  const velocity = this.storage.velocityStorage.element(boidIndex);
  const speed = length(velocity);
  
  // Normalize speed for color mixing (assuming max speed around 10)
  const normalizedSpeed = speed.div(10.0).saturate();
  
  // Mix between two colors based on speed
  const color = mix(
    vec3(this.config.colorA.r, this.config.colorA.g, this.config.colorA.b),
    vec3(this.config.colorB.r, this.config.colorB.g, this.config.colorB.b),
    normalizedSpeed
  );
  
  return vec4(color, 1.0);
});

1. Identify the Boid: Just like the vertex shader, it uses instanceIndex to know which boid it's coloring.
2. Fetch Velocity: It reads the velocity for that specific boid from the velocityStorage buffer.
3. Calculate Speed: It calculates the magnitude (or length) of the velocity vector to get the boid's scalar speed.
4. Normalize and Interpolate: The speed is normalized to a range of 0.0 to 1.0 (the .saturate() call clamps the value). This normalizedSpeed is then used as the mixing factor in TSL's mix() function. This function performs a linear interpolation between two colors. A boid with a speed of 0 will be one color, a boid at maximum speed will be the other, and boids in between will be a blended color.
5. Output Color: The final calculated color is returned as a vec4 (with an alpha value of 1.0 for full opacity).

Summary of the Data Flow per Frame

The entire process is a highly optimized loop that happens on every frame:

1. The CPU tells the GPU to run the boids simulation compute shaders.
2. The compute shaders run, updating the positions and velocities within the GPU's storage buffers.
3. The CPU tells the GPU to render the InstancedMesh.
4. The GPU executes the vertex shader for each boid, reading the *newly updated* positions and velocities to calculate where each boid should be drawn and how it should be oriented.
5. The GPU then executes the fragment shader for each pixel of each boid, reading the velocities to determine the correct color.

This entire process minimizes CPU involvement and avoids costly data transfers between the CPU and GPU, enabling a smooth, real-time visualization of thousands of agents.

Complete Boids Visualization Code

Here is the complete code for the BoidsVisualization class, which handles all aspects of rendering the boids flock using Three.js and TSL shaders.

import * as THREE from 'three/webgpu';
import { 
  Fn, 
  normalize, 
  length, 
  mix,
  vec3,
  vec4,
  instanceIndex,
  positionLocal,
  cameraProjectionMatrix,
  cameraViewMatrix,
  mat3,
  cross
} from 'three/tsl';
import { BoidsSimulation } from './boids';

export interface BoidsVisualizationConfig {
  particleSize: number;
  colorA: THREE.Color;
  colorB: THREE.Color;
  useTriangles: boolean; // true for triangles, false for quads
}

export class BoidsVisualization {
  private scene: THREE.Scene;
  private camera!: THREE.PerspectiveCamera;
  private mesh!: THREE.InstancedMesh;
  private material!: THREE.NodeMaterial;
  private config: BoidsVisualizationConfig;
  private storage: any; // Use any type to avoid import issues
  private count: number;

  constructor(
    simulation: BoidsSimulation,
    config: Partial<BoidsVisualizationConfig> = {}
  ) {
    this.config = {
      particleSize: 1.0,
      colorA: new THREE.Color(0x00ff00),
      colorB: new THREE.Color(0xff0000),
      useTriangles: true,
      ...config
    };

    this.storage = simulation.getStorage();
    this.count = simulation.getConfig().count;

    this.scene = new THREE.Scene();
    this.setupCamera();
    this.setupGeometry();
    this.setupMaterial();
    this.setupMesh();
  }

  private setupCamera(): void {
    this.camera = new THREE.PerspectiveCamera(
      50, 
      window.innerWidth / window.innerHeight, 
      1, 
      10000
    );
    this.camera.position.set(0, 0, 1600);
  }

  private setupGeometry(): THREE.BufferGeometry {
    if (this.config.useTriangles) {
      // Simple triangle pointing upward (in local space)
      const geometry = new THREE.BufferGeometry();
      const vertices = new Float32Array([
        0.0,  0.5, 0.0,  // top vertex
       -0.1, -0.5, 0.0,  // bottom left
        0.1, -0.5, 0.0   // bottom right
      ]);
      geometry.setAttribute('position', new THREE.BufferAttribute(vertices, 3));
      return geometry;
    } else {
      // Simple quad as fallback
      return new THREE.PlaneGeometry(0.6, 0.6);
    }
  }

  private setupMaterial(): void {
    this.material = new THREE.NodeMaterial();
    
    // Create vertex shader that positions and orients each boid
    const boidVertexShader = Fn(() => {
      // Get the current instance index (which boid we're rendering)
      const boidIndex = instanceIndex;
      
      // Get position and velocity from compute storage
      const boidPosition = this.storage.positionStorage.element(boidIndex);
      const boidVelocity = this.storage.velocityStorage.element(boidIndex);
      
      // Transform local vertex position
      const localPos = positionLocal.toVar();
      localPos.mulAssign(this.config.particleSize);
      
      // Create a rotation matrix to align the boid with its velocity
      const velocity = normalize(boidVelocity.add(vec3(0.001, 0.001, 0.001))); // Add epsilon to avoid zero velocity

      const forward = velocity.toVar('forward');
      const up = vec3(0.0, 1.0, 0.0).toVar('up');
      const right = normalize(cross(up, forward)).toVar('right');
      const newUp = normalize(cross(forward, right)).toVar('newUp');
      
      const rotationMatrix = mat3(
        right.x, forward.x, newUp.x,
        right.y, forward.y, newUp.y,
        right.z, forward.z, newUp.z
      ).toVar();
      
      const rotatedPos = rotationMatrix.mul(localPos);
      
      // Translate to boid position
      const worldPos = rotatedPos.add(boidPosition);
      
      // Transform to clip space
      return cameraProjectionMatrix.mul(cameraViewMatrix).mul(vec4(worldPos, 1.0));
    });

    // Create fragment shader for coloring based on speed
    const boidFragmentShader = Fn(() => {
      const boidIndex = instanceIndex;
      const velocity = this.storage.velocityStorage.element(boidIndex);
      const speed = length(velocity);
      
      // Normalize speed for color mixing (assuming max speed around 10)
      const normalizedSpeed = speed.div(10.0).saturate();
      
      // Mix between two colors based on speed
      const color = mix(
        vec3(this.config.colorA.r, this.config.colorA.g, this.config.colorA.b),
        vec3(this.config.colorB.r, this.config.colorB.g, this.config.colorB.b),
        normalizedSpeed
      );
      
      return vec4(color, 1.0);
    });

    // Assign shaders to material
    this.material.vertexNode = boidVertexShader();
    this.material.colorNode = boidFragmentShader();
    this.material.side = THREE.DoubleSide;
    this.material.transparent = true;
  }

  private setupMesh(): void {
    const geometry = this.setupGeometry();
    this.mesh = new THREE.InstancedMesh(geometry, this.material, this.count);
    this.mesh.frustumCulled = false; // Disable frustum culling for performance
    this.scene.add(this.mesh);
  }

  public render(renderer: THREE.WebGPURenderer): void {
    renderer.render(this.scene, this.camera);
  }

  public getScene(): THREE.Scene {
    return this.scene;
  }

  public getCamera(): THREE.PerspectiveCamera {
    return this.camera;
  }

  public updateConfig(config: Partial<BoidsVisualizationConfig>): void {
    Object.assign(this.config, config);
    
    // Recreate material if colors changed
    if (config.colorA || config.colorB) {
      this.setupMaterial();
      this.mesh.material = this.material;
    }
    
    // Recreate geometry if triangle/quad setting changed
    if (config.useTriangles !== undefined) {
      const newGeometry = this.setupGeometry();
      this.mesh.geometry.dispose();
      this.mesh.geometry = newGeometry;
    }
  }

  public onWindowResize(width: number, height: number): void {
    this.camera.aspect = width / height;
    this.camera.updateProjectionMatrix();
  }

  public dispose(): void {
    this.mesh.geometry.dispose();
    this.material.dispose();
    this.scene.clear();
  }
}