Starting point for reduce

reduce.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 {
        T : i32, // number of threads launched
        P : i32, // side size of square grid (P*P >= T)
        S : i32, // step
      }
      
      @group(0) @binding(0) var<uniform>             param: Params;
      @group(0) @binding(1) var<storage, read_write> E_in  : array<u32>;
      @group(0) @binding(2) var<storage, read_write> E_out : array<u32>;

      @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 id  = gid.x + gid.y * param.P;
          if (id < param.T) {
            // active thread, otherwise padding
            
            // T[id] = ...
          }
      }
    `,
  });

  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',
        },
      },
      {
        binding: 2,
        visibility: GPUShaderStage.COMPUTE,
        buffer: {
          type: 'storage',
        },
      }
    ],
  })

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

  E = 317;
 
  // GPU read/write buffer
  const bufferA = device.createBuffer({
    label: 'A',
    size: E * 4 /*u32*/,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
  });
  // place some initial data inside
 	const input = new Uint32Array(E);
  for (var n=0;n < E;++n) {
    input[n] = 1;
  }
  // copy our init data from CPU to GPU
  device.queue.writeBuffer(bufferA, 0, input); 
  // allocate a second buffer
  const bufferB = device.createBuffer({
    label: 'B',
    size: E * 4 /*u32*/,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
  });

  // create a buffer on the GPU to get a copy of the results
  const resultBuffer = device.createBuffer({
    label: 'result buffer',
    size: 4 /*a single u32*/,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
  });
  
  // create a buffer for the uniform parameters
  const params = device.createBuffer({
    label: 'uniform buffer',
    size: 3 * 4 /*u32*/,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  console.log(params)

  var iter = 0;

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

  var n_steps = 1; // ... TODO
  for (var step=0 ; step < n_steps ; step++) {

    var T = E; // ... TODO
    var P = Math.ceil(Math.sqrt(T));
    console.log('T = ',T,' P = ',P, ' P*P = ',P*P);
    if (P*P < T) {
      console.error('incorrect grid size');
    }

    // update uniform parameters
    device.queue.writeBuffer(params,0,new Uint32Array([T,P,step]));
    
    // 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 /*TODO*/, resource: { buffer: bufferA}, },
        { binding: 2 /*TODO*/, resource: { buffer: bufferB}, },
      ],
    });
    // Encode commands to do the computation
    const encoder = device.createCommandEncoder();
    const pass = encoder.beginComputePass();  
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(P,P);
    pass.end();
    // Finish encoding and submit the commands
    const commandBuffer = encoder.finish();
    // Submit commands for this step
    device.queue.submit([commandBuffer]);

  }

  // Encode a command to copy the results to a mappable buffer.
  const encoder = device.createCommandEncoder();
  encoder.copyBufferToBuffer(bufferB /*TODO*/ , 0, resultBuffer, 0, resultBuffer.size);
  // Finish encoding and submit the commands
  const commandBuffer = encoder.finish();
  // Submit commands for this step
  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);
  
}

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

main();