export function runWorkgroupVariableTest()

in src/webgpu/shader/execution/expression/call/builtin/atomics/harness.ts [114:206]
54 lines of code
3 McCabe index (conditional complexity)

export function runWorkgroupVariableTest({
  t,
  workgroupSize, // Workgroup X-size
  dispatchSize, // Dispatch X-size
  wgNumElements, // Number of 32-bit elements in 'wg' array. Output buffer is sized to wgNumElements * dispatchSize.
  initValue, // 32-bit initial value used to fill 'wg' array
  // Atomic op source executed by the compute shader, NOTE: 'id' is local_invocation_index,
  // `wg` is a workgroup array of atomics of size `workgroupSize`, `output` is a storage array of non-atomics of size
  // `workgroupSize * dispatcSize` to which each dispatch of `wg` gets copied to (dispatch 0 to first workgroupSize elements,
  // dispatch 1 to second workgroupSize elements, etc.).
  op,
  expected, // Expected values array to compare against output buffer
  extra, // Optional extra WGSL source
}: {
  t: GPUTest;
  workgroupSize: number;
  dispatchSize: number;
  wgNumElements: number;
  initValue: number;
  op: string;
  expected: TypedArrayBufferView;
  extra?: string;
}) {
  assert(expected.length === wgNumElements * dispatchSize, "'expected' buffer size is incorrect");

  const scalarType = expected instanceof Uint32Array ? 'u32' : 'i32';
  const arrayType = typedArrayCtor(scalarType);

  const wgsl = `
    var<workgroup> wg: array<atomic<${scalarType}>, ${wgNumElements}>;

    // Result of each workgroup is written to output[workgroup_id.x]
    @group(0) @binding(0)
    var<storage, read_write> output: array<${scalarType}, ${wgNumElements * dispatchSize}>;

    @compute @workgroup_size(${workgroupSize})
    fn main(
        @builtin(local_invocation_index) local_invocation_index: u32,
        @builtin(workgroup_id) workgroup_id : vec3<u32>
        ) {
      let id = ${scalarType}(local_invocation_index);
      let global_id = ${scalarType}(workgroup_id.x * ${wgNumElements} + local_invocation_index);

      // Initialize workgroup array
      if (local_invocation_index == 0) {
        for (var i = 0u; i < ${wgNumElements}; i++) {
          atomicStore(&wg[i], bitcast<${scalarType}>(${initValue}u));
        }
      }
      workgroupBarrier();

      ${op};

      // Copy results to output buffer
      workgroupBarrier();
      if (local_invocation_index == 0) {
        for (var i = 0u; i < ${wgNumElements}; i++) {
          output[(workgroup_id.x * ${wgNumElements}) + i] = atomicLoad(&wg[i]);
        }
      }
    }
    ${extra || ''}
    `;

  const pipeline = t.device.createComputePipeline({
    layout: 'auto',
    compute: {
      module: t.device.createShaderModule({ code: wgsl }),
      entryPoint: 'main',
    },
  });

  const outputBuffer = t.createBufferTracked({
    size: wgNumElements * dispatchSize * arrayType.BYTES_PER_ELEMENT,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
  });

  const bindGroup = t.device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [{ binding: 0, resource: { buffer: outputBuffer } }],
  });

  // Run the shader.
  const encoder = t.device.createCommandEncoder({ label: 'runWorkgroupVariableTest' });
  const pass = encoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(dispatchSize);
  pass.end();
  t.queue.submit([encoder.finish()]);

  t.expectGPUBufferValuesEqual(outputBuffer, expected);
}