WebGPU twoSum

README.md

WebGPU twoSum

view on jsgist

index.css

@import url(https://webgpufundamentals.org/webgpu/resources/webgpu-lesson.css);

table { border-collapse: collapse; }
td { border: 1px solid #888; padding: 0.25em }
td:empty { background-color: #888; }

index.html

<p>C++ version is <a href="https://onlinegdb.com/29Tbs-_rR">here</a>

<p>GPU-b should match js-b</p> 

index.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: 'doubling compute module',
    code: `
      struct float2 {
        x: f32,
        y: f32
      }

      fn twoSum(a: f32, b: f32) -> float2 {
        let x = a + b;
        let bv = x - a;
        let av = x - bv;
        let br = b - bv;
        let ar = a - av;
        return float2(x, ar + br);
      }

      @group(0) @binding(0) var<storage, read> input: array<f32>;
      @group(0) @binding(1) var<storage, read_write> output: array<float2>;

      @compute @workgroup_size(1) fn computeSomething(
        @builtin(global_invocation_id) id: vec3<u32>
      ) {
        let outNdx = id.x;
        let inNdx = outNdx * 2;
        output[outNdx] = twoSum(input[inNdx], input[inNdx + 1]);
      }
    `,
  });

  const pipeline = device.createComputePipeline({
    label: 'twoSum compute pipeline',
    layout: 'auto',
    compute: {
      module,
      entryPoint: 'computeSomething',
    },
  });

  const data = [
    1, 3,
    5, 6,
    1e10, 1e-10,
  ];
  const input = new Float32Array(data);

  // create a buffer on the GPU to hold our input
  const inputBuffer = device.createBuffer({
    label: 'input buffer',
    size: input.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
  });
  // Copy our input data to that buffer
  device.queue.writeBuffer(inputBuffer, 0, input);

  // create a buffer on the GPU to hold our output
  const outputBuffer = device.createBuffer({
    label: 'output buffer',
    size: input.byteLength * 2,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
  });

  // create a buffer on the GPU to get a copy of the results
  const resultBuffer = device.createBuffer({
    label: 'result buffer',
    size: outputBuffer.size,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
  });

  // Setup a bindGroup to tell the shader which
  // buffers to use for the computation
  const bindGroup = device.createBindGroup({
    label: 'bindGroup for work buffer',
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: inputBuffer } },
      { binding: 1, resource: { buffer: outputBuffer } },
    ],
  });

  // Encode commands to do the computation
  const encoder = device.createCommandEncoder({
    label: 'twoSum encoder',
  });
  const pass = encoder.beginComputePass({
    label: 'twoSum compute pass',
  });
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(input.length);
  pass.end();

  // Encode a command to copy the results to a mappable buffer.
  encoder.copyBufferToBuffer(outputBuffer, 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 Float32Array(resultBuffer.getMappedRange().slice());
  resultBuffer.unmap();

  const x = new Float32Array(1);
  const bv = new Float32Array(1);
  const av = new Float32Array(1);
  const br = new Float32Array(1);
  const ar = new Float32Array(1);
  const a = new Float32Array(1);
  const b = new Float32Array(1);

  function twoSum(_a, _b) {
    a[0] = _a;
    b[0] = _b;
    x[0] = a[0] + b[0];
    bv[0] = x[0] - a[0];
    av[0] = x[0] - bv[0];
    br[0] = b[0] - bv[0];
    ar[0] = a[0] - av[0];   
    return [x[0], ar[0] + br[0]]; 
  }

  const addRow = makeTable(document.body, ['a', 'b', '', 'js-a', 'js-b', 'jsSum', '', 'GPU-a', 'GPU-b', 'GPU-sum']);

  for (let i = 0; i < data.length; i += 2) {
    const a = data[i];
    const b = data[i + 1];
    const js = twoSum(a, b);
    addRow([a, b, '', js[0], js[1], js[0] + js[1], '', result[i], result[i + 1], result[i] + result[i + 1]].map(v => v.toExponential?.call(v)));
  }
}

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

function c(tag, children = [], text = '') {
  const e = document.createElement(tag);
  e.textContent = text;
  children.forEach(c => e.appendChild(c));
  return e;
}
  
function makeTable(parent, columnNames) {
  const makeRow = (arr, tag = 'td') => c('tr', arr.map(v => c(tag, [], v)));

  const tbody = c('tbody');
  parent.appendChild(c('table', [
    c('thead', [makeRow(columnNames, 'th')]),
    tbody,
  ]));
  return function(columnData) {
    tbody.appendChild(makeRow(columnData));
  };
}

main();

console.log('1e-10 through Float32Array = ', new Float32Array([1e-10])[0].toExponential())

jsGist.json

{"name":"WebGPU twoSum","settings":{},"filenames":["index.html","index.css","index.js"]}