Iterations

class_4_gof.js

async function main()
{
  const adapter = await navigator.gpu?.requestAdapter();
  const device = await adapter?.requestDevice();
  if (!device) {
    fail('need a browser that supports WebGPU');
    return;
  }

  const module = device.createShaderModule({
    label: 'compute module',
    code: `
      struct Params {
        N : i32,
        K : i32,
      }
      
      @group(0) @binding(0) var<uniform>             param: Params;
      @group(0) @binding(1) var<storage, read_write> T    : array<u32>;

      fn gaussian(uv : vec2<f32>, dev : f32) -> f32
      {
        let dsq   = dot(uv,uv);          
        let g     = 2.0 * exp(- dsq * dev) / exp(1.0);
        return g;
      }
      
      fn u32_to_color(rgba : u32) -> vec3<f32>
      {
        return vec3<f32>(
          f32( rgba      & 255),
          f32((rgba>> 8) & 255),
          f32((rgba>>16) & 255)
        );
      }

      @compute @workgroup_size(1,1) fn computeSomething(
        @builtin(global_invocation_id) u_gid: vec3<u32>,
        @builtin(local_invocation_id)  u_lid: vec3<u32>
      ) {
          let gid = vec3<i32>(u_gid);
          let x   = f32(gid.x) / f32(param.N);
          let y   = f32(gid.y) / f32(param.N);
          let id  = gid.x + gid.y * param.N;

          var argb = vec3<f32>(0.0);
          var ak   = 0.0;
          for (var j =   gid.y - param.K ;
                   j <=  gid.y + param.K ; j ++) {
            for (var i =   gid.x - param.K ;
                     i <=  gid.x + param.K ; i ++) {
                if (i >= 0 && j >= 0 
                   && i < param.N && j < param.N) 
                {
                   // i,j on valid coordinate
                   // -> read pixel value
                   let pid  = i + j * param.N;
                   let prgb = u32_to_color(T[pid]);
                   // -> sample kernel
                   var pk  = gaussian(vec2<f32>(
                                    f32(i - gid.x) * 0.5 / f32(param.K),
                                    f32(j - gid.y) * 0.5 / f32(param.K)
                                    
                                    ), 3.0 + x*x*x * 50.0);
                   // -> cumulate
                   argb = argb + prgb * pk;
                   ak   = ak + pk;
                }
            }
          }
          argb    = argb / ak;
          let rgb = vec3<u32>( argb );
          T[id]   = rgb.x | (rgb.y << 8) | (rgb.z << 16);
   
/*          
          // visualize gaussian kernel
          let delta = vec2<f32>(x,y) - vec2<f32>(0.5,0.5);          
          let g     = gaussian(delta, 9.0);
          T[id]     =  u32(g*255.0)
                    | (u32(g*255.0) << 8) 
                    | (u32(g*255.0) <<16); 
*/          
      }
    `,
  });

  const piplayout_group0 = device.createBindGroupLayout({
    label: 'piplayout_group0',
    entries: [
      {
        binding: 0,
        visibility: GPUShaderStage.COMPUTE,
        buffer: {
          type: 'uniform',
        },
      },
      {
        binding: 1,
        visibility: GPUShaderStage.COMPUTE,
        buffer: {
          type: 'storage',
        },
      },
    ],
  })

  const pipeline = device.createComputePipeline({
    label: 'compute pipeline',
    layout: device.createPipelineLayout({
        bindGroupLayouts: [piplayout_group0],
    }),
    compute: {
      module,
      entryPoint: 'computeSomething',
    },
  });

  N = 256;
  K = 16;
 
  var img = document.querySelector('#myimage'); 
  var canvas = document.createElement('canvas'); 
  var ctx = canvas.getContext('2d'); 
  canvas.width  = img.width; 
  canvas.height = img.height;  
  ctx.drawImage(img, 0, 0);  
  var rgba = ctx.getImageData(0, 0, img.width, img.height).data;
  console.log(rgba);
 
  // GPU read/write buffer
  const bufferT = device.createBuffer({
    label: 'T',
    size: N * N * 4 /*u32*/,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
  });
  // place some initial data inside
 	const input = new Uint32Array(N*N);
  for (var n=0;n < N*N;++n) {
    input[n] =  rgba[n*4 + 0]
             | (rgba[n*4 + 1] <<  8)
             | (rgba[n*4 + 2] << 16)
             ;
  }
  // copy our init data from CPU to GPU
  device.queue.writeBuffer(bufferT, 0, input); 
  
  // create a buffer on the GPU to get a copy of the results
  const resultBuffer = device.createBuffer({
    label: 'result buffer',
    size: N*N*4 /*u32*/,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
  });
  
  // create a buffer for the uniform parameters
  const params = device.createBuffer({
    label: 'uniform buffer',
    size: 2 * 4 /*u32*/,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  console.log(params)
  device.queue.writeBuffer(params,0,new Uint32Array([N,K]));

  var iter = 0;

  console.log('==== MAIN LOOP ====');

  // create an offscreen canvas
  var canvas = document.createElement('canvas');
  var ctx = canvas.getContext('2d');
  // size the canvas to your desired image
  canvas.width  = N;
  canvas.height = N;
  // create a new img object
  var image = new Image();
  // set the img.src to the canvas data url
  image.src = canvas.toDataURL();
  // append the new img object to the page
  document.body.appendChild(image);

  setInterval( async () => {

    // Setup a bindGroup to tell the shader which
    // buffer to use for the computation
    const bindGroup = device.createBindGroup({
      label: 'bindGroup',
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        { binding: 0, resource: { buffer: params }, },
        { binding: 1, resource: { buffer: bufferT}, },
      ],
    });
    // Encode commands to do the computation
    const encoder = device.createCommandEncoder({
      label: 'encoder',
    });
    const pass = encoder.beginComputePass({
      label: 'compute pass',
    });  
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(N,N);
    pass.end();
    // Encode a command to copy the results to a mappable buffer.
    encoder.copyBufferToBuffer(bufferT, 0, resultBuffer, 0, resultBuffer.size);
    // Finish encoding and submit the commands
    const commandBuffer = encoder.finish();

    device.queue.submit([commandBuffer]);

    // Read the results
    await resultBuffer.mapAsync(GPUMapMode.READ);
    const result = new Uint32Array(resultBuffer.getMappedRange().slice());
    resultBuffer.unmap();

    // console.log('result', result);
    console.log('iteration ', iter++);
    
    // get the imageData and pixel array from the canvas
    var imgData = ctx.getImageData(0, 0, N, N);
    var data = imgData.data;
    // manipulate some pixel elements
    for (var i = 0; i < data.length; i++) {
      data[i*4 + 0] = (result[i]     &255);
      data[i*4 + 1] = (result[i]>>8) &255;
      data[i*4 + 2] = (result[i]>>16)&255;
      data[i*4 + 3] = 255; // make this pixel opaque
    }
    // put the modified pixels back on the canvas
    ctx.putImageData(imgData, 0, 0);
    image.src = canvas.toDataURL();

  }
  , 1000 /*delay*/ )
  
}

function fail(msg) {
  // eslint-disable-next-line no-alert
  alert(msg);
}

main();