Fire Effect (Animated)

<?php
let div = document.createElement('div');
document.body.appendChild( div );
div.style['font-size'] = '20pt';
function log( s )
{
  console.log( s );
  let args = [...arguments].join(' ');
  div.innerHTML += args + '<br><br>';
}

log('WebGPU Compute Example');


if (!navigator.gpu) { log("WebGPU is not supported (or is it disabled? flags/settings)"); return; }

const adapter = await navigator.gpu.requestAdapter();
const device  = await adapter.requestDevice();

const imgWidth = 256;
const imgHeight = imgWidth;

const bufferSize = imgWidth * imgHeight;

// ----------------------------------------------------------

// Basic canvas which will be used to display the output from the compute shader

let canvasa = document.createElement('canvas');
document.body.appendChild( canvasa ); canvasa.height = canvasa.width = imgWidth;
const context = canvasa.getContext('webgpu');
const presentationFormat = navigator.gpu.getPreferredCanvasFormat(); 
console.log('presentationFormat:', presentationFormat );

context.configure({ device: device, 
                    usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC | GPUTextureUsage.COPY_DST,
                    format: "rgba8unorm" /*presentationFormat*/  });

let canvasTexture = context.getCurrentTexture();

// ----------------------------------------------------------

const texture1 = device.createTexture({
  size: [imgWidth, imgHeight, 1],
  format: "rgba8unorm",
  usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.COPY_SRC | GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.STORAGE_BINDING
});

// ----------------------------------------------------------

const timerUniformBuffer = device.createBuffer({
  size: 4, 
  usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
});

const timestep  = new Float32Array( [0.0] );

device.queue.writeBuffer(timerUniformBuffer,   0, timestep             );

// ----------------------------------------------------------

function Colour(r, g, b) {
    this.r = r;
    this.g = g;
    this.b = b;
    this.a = 255;
}

let colors = [];

for (let i = 0; i < 256; i++) {
  colors[i] = new Colour(0, 0, 0);
}

for (let i = 0; i < 32; ++i) {
  // black to blue
  colors[i].b = i << 1;

  // blue to red
  colors[i + 32].r = i << 3;
  colors[i + 32].b = 64 - (i << 1);

  // red to yellow
  colors[i + 64].r = 255;
  colors[i + 64].g = i << 3;

  // yellow to white
  colors[i + 96].r = 255;
  colors[i + 96].g = 255;
  colors[i + 96].b = i << 2;
  colors[i + 128].r = 255;
  colors[i + 128].g = 255;
  colors[i + 128].b = 64 + (i << 2);
  colors[i + 160].r = 255;
  colors[i + 160].g = 255;
  colors[i + 160].b = 128 + (i << 2);
  colors[i + 192].r = 255;
  colors[i + 192].g = 255;
  colors[i + 192].b = 192 + i;
  colors[i + 224].r = 255;
  colors[i + 224].g = 255;
  colors[i + 224].b = 224 + i;
}

colors = colors.map( (a)=>{ return Object.values(a); } ).flat();

const colorArray = new Float32Array( colors ); // vec4<f32>

const colorBuffer = device.createBuffer({ size: colorArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST  | GPUBufferUsage.COPY_SRC });
device.queue.writeBuffer(colorBuffer, 0, colorArray);

// ----------------------------------------------------------

const bufArray0 = new Float32Array( imgWidth * imgHeight ); 

const bufBuffer0 = device.createBuffer({ size: bufArray0.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST  | GPUBufferUsage.COPY_SRC });
device.queue.writeBuffer(bufBuffer0, 0, colorArray);


const bufBuffer1 = device.createBuffer({ size: bufArray0.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST  | GPUBufferUsage.COPY_SRC });
device.queue.writeBuffer(bufBuffer1, 0, colorArray);

// ----------------------------------------------------------

const GCOMPUTE = GPUShaderStage.COMPUTE;

// Bind group layout and bind group
const bindGroupLayout = device.createBindGroupLayout({
  entries: [ {binding: 0, visibility: GCOMPUTE, buffer:  { type: "uniform"  }   },
             {binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: {type: "read-only-storage"}  },
             {binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"}  },
             {binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"}  },
             {binding: 4, visibility: GCOMPUTE, storageTexture: {format:"rgba8unorm", access:"write-only", viewDimension:"2d"}   }
           ]
});

const bindGroup0 = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [  {   binding: 0,  resource: { buffer: timerUniformBuffer  } },
                  {   binding: 1,  resource: { buffer: colorBuffer         } },
                {   binding: 2,  resource: { buffer: bufBuffer0          } }, // swap buf0/1
                {   binding: 3,  resource: { buffer: bufBuffer1          } },
                {   binding: 4,  resource: texture1.createView()           }
    ]
});

const bindGroup1 = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [  {   binding: 0,  resource: { buffer: timerUniformBuffer  } },
                  {   binding: 1,  resource: { buffer: colorBuffer         } },
                {   binding: 2,  resource: { buffer: bufBuffer1          } }, // swap buf0/1
                {   binding: 3,  resource: { buffer: bufBuffer0          } },
                {   binding: 4,  resource: texture1.createView()           }
    ]
});

  
// Compute shader code
const computeShader = ` 
@group(0) @binding(0) var<uniform>             mytimer  : f32; // timer increments each frame
@group(0) @binding(1) var<storage, read>       mycolors : array< vec4<f32>, 256 >;
@group(0) @binding(2) var<storage, read_write> buf0     : array< f32, ${bufArray0.length} >; // width x height
@group(0) @binding(3) var<storage, read_write> buf1     : array< f32, ${bufArray0.length} >; // width x height
@group(0) @binding(4) var myTexture1: texture_storage_2d<rgba8unorm, write>; // output image

const imgWidth  = u32( ${imgWidth}  ); 
const imgHeight = u32( ${imgHeight} ); 
//const imgSize = vec2<f32>(imgWidth, imgHeight);
    
// Simple random number generator based on the index
fn rand(seed: u32) -> f32 {
    var x = seed * 1234567u + 987654321u;
    x = (x >> 16u) ^ x;
    x = x * 1103515245u + 12345u;
    return f32(x & 0x7fffffffu) / f32(0x7fffffffu);
}

fn random(uv: vec2<f32>) -> f32 {
    return fract(sin(dot(uv, vec2<f32>(12.9898, 78.233))) * 43758.5453);
}

@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) globalId      : vec3<u32>,
        @builtin(local_invocation_id)  localId       : vec3<u32>,
        @builtin(workgroup_id)         workgroupId   : vec3<u32>,
        @builtin(num_workgroups)       workgroupSize : vec3<u32>
        ) 
{
    var coords = vec2<u32>(globalId.xy);

    if (coords.x >= imgWidth || coords.y >= imgHeight) {
        return;
    }

    let index = coords.y * imgWidth + coords.x;
    
    // Simulate fire effect
    if (coords.y == imgHeight - 1) 
    {
        // Randomize the bottom row of the fire array
        if ( random( vec2<f32>(f32(index)/1024.0, mytimer*0.1 ) ) > 0.5) {
            buf1[index] = 255.0;
        } else {
            buf1[index] = 0;
        }
    } 
    else 
    {
        // Move fire upwards
        var temp = f32(buf0[index]);
        temp    += f32(buf0[index + 1]);
        temp    += f32(buf0[index - 1]);
        temp    += f32(buf0[index + imgWidth]);
        temp /= 3.975;

        if (temp > 1.0) {
            temp -= 1.0;
        }

        buf1[index] = f32(clamp(temp, 0.0, 255.0));
    }

    var finalcolor = mycolors[ u32( buf0[ index ] ) ];
    
    finalcolor *= 1.0/255;
    finalcolor.w = 1.0;
    
    // Store the result in the output texture
    textureStore(myTexture1, vec2<i32>(globalId.xy), finalcolor);
}
`;
  

// Pipeline setup
const computePipeline = device.createComputePipeline({
    layout :   device.createPipelineLayout({bindGroupLayouts: [bindGroupLayout]}),
    compute: { module    : device.createShaderModule({code:computeShader}),
               entryPoint: "main" }
});


async function frame()
{
  timestep[0] = timestep[0] + 0.01;
  device.queue.writeBuffer(timerUniformBuffer,   0, timestep             );


  // Commands submission
  const commandEncoder = device.createCommandEncoder();

  {
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup0);
  // workgroup size of 8x8 on the wgsl shader
  passEncoder.dispatchWorkgroups( imgWidth/8, imgWidth/8 );
  await passEncoder.end();
  }

  {
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup1);
  // workgroup size of 8x8 on the wgsl shader
  passEncoder.dispatchWorkgroups( imgWidth/8, imgWidth/8 );
  await passEncoder.end();
  }

  canvasTexture = context.getCurrentTexture();

  await
  commandEncoder.copyTextureToTexture( { texture: texture1 },
                                      { texture: canvasTexture },
                                      { width:imgWidth, height:imgHeight, depthOrArrayLayers:1} );

  // Submit GPU commands.
  const gpuCommands = commandEncoder.finish();
  await device.queue.submit([gpuCommands]);
  
  requestAnimationFrame(frame);
}

frame();
Complete Code

Things to Try

Resources and Links
