export const description = `Test compute shader builtin variables`;

import { makeTestGroup } from '../../../../common/framework/test_group.js';
import { iterRange } from '../../../../common/util/util.js';
import { AllFeaturesMaxLimitsGPUTest } from '../../../gpu_test.js';

export const g = makeTestGroup(AllFeaturesMaxLimitsGPUTest);

// Test that the values for each input builtin are correct.
g.test('inputs')
  .desc(`Test compute shader builtin inputs values`)
  .params(u =>
    u
      .combine('method', ['param', 'struct', 'mixed'] as const)
      .combine('dispatch', ['direct', 'indirect'] as const)
      .combineWithParams([
        {
          groupSize: { x: 1, y: 1, z: 1 },
          numGroups: { x: 1, y: 1, z: 1 },
        },
        {
          groupSize: { x: 8, y: 4, z: 2 },
          numGroups: { x: 1, y: 1, z: 1 },
        },
        {
          groupSize: { x: 1, y: 1, z: 1 },
          numGroups: { x: 8, y: 4, z: 2 },
        },
        {
          groupSize: { x: 3, y: 7, z: 5 },
          numGroups: { x: 13, y: 9, z: 11 },
        },
      ] as const)
      .beginSubcases()
  )
  .fn(t => {
    const linear_indexing = t.hasLanguageFeature('linear_indexing');
    const invocationsPerGroup = t.params.groupSize.x * t.params.groupSize.y * t.params.groupSize.z;
    const totalInvocations =
      invocationsPerGroup * t.params.numGroups.x * t.params.numGroups.y * t.params.numGroups.z;

    // Generate the structures, parameters, and builtin expressions used in the shader.
    let params = '';
    let structures = '';
    let local_id = '';
    let local_index = '';
    let global_id = '';
    let group_id = '';
    let num_groups = '';
    let global_index = '';
    let group_index = '';
    switch (t.params.method) {
      case 'param':
        params = `
          @builtin(local_invocation_id) local_id : vec3<u32>,
          @builtin(local_invocation_index) local_index : u32,
          @builtin(global_invocation_id) global_id : vec3<u32>,
          @builtin(workgroup_id) group_id : vec3<u32>,
          @builtin(num_workgroups) num_groups : vec3<u32>,
          ${linear_indexing ? '@builtin(global_invocation_index) global_index : u32,' : ''}
          ${linear_indexing ? '@builtin(workgroup_index) group_index : u32,' : ''}
        `;
        local_id = 'local_id';
        local_index = 'local_index';
        global_id = 'global_id';
        group_id = 'group_id';
        num_groups = 'num_groups';
        if (linear_indexing) {
          global_index = 'global_index';
          group_index = 'group_index';
        }
        break;
      case 'struct':
        structures = `struct Inputs {
            @builtin(local_invocation_id) local_id : vec3<u32>,
            @builtin(local_invocation_index) local_index : u32,
            @builtin(global_invocation_id) global_id : vec3<u32>,
            @builtin(workgroup_id) group_id : vec3<u32>,
            @builtin(num_workgroups) num_groups : vec3<u32>,
            ${linear_indexing ? '@builtin(global_invocation_index) global_index : u32,' : ''}
            ${linear_indexing ? '@builtin(workgroup_index) group_index : u32,' : ''}
          };`;
        params = `inputs : Inputs`;
        local_id = 'inputs.local_id';
        local_index = 'inputs.local_index';
        global_id = 'inputs.global_id';
        group_id = 'inputs.group_id';
        num_groups = 'inputs.num_groups';
        if (linear_indexing) {
          global_index = 'inputs.global_index';
          group_index = 'inputs.group_index';
        }
        break;
      case 'mixed':
        structures = `struct InputsA {
          @builtin(local_invocation_index) local_index : u32,
          @builtin(global_invocation_id) global_id : vec3<u32>,
        };
        struct InputsB {
          @builtin(workgroup_id) group_id : vec3<u32>
        };`;
        params = `@builtin(local_invocation_id) local_id : vec3<u32>,
                  inputsA : InputsA,
                  inputsB : InputsB,
                  @builtin(num_workgroups) num_groups : vec3<u32>,
                  ${linear_indexing ? '@builtin(global_invocation_index) global_index : u32,' : ''}
                  ${linear_indexing ? '@builtin(workgroup_index) group_index : u32,' : ''}
                  `;
        local_id = 'local_id';
        local_index = 'inputsA.local_index';
        global_id = 'inputsA.global_id';
        group_id = 'inputsB.group_id';
        num_groups = 'num_groups';
        if (linear_indexing) {
          global_index = 'global_index';
          group_index = 'group_index';
        }
        break;
    }

    // WGSL shader that stores every builtin value to a buffer, for every invocation in the grid.
    const wgsl = `
      struct Outputs {
        local_id: vec3u,
        local_index: u32,
        global_id: vec3u,
        group_id: vec3u,
        num_groups: vec3u,
        ${linear_indexing ? 'global_index : u32,' : ''}
        ${linear_indexing ? 'group_index : u32,' : ''}
      };
      @group(0) @binding(0) var<storage, read_write> outputs : array<Outputs>;

      ${structures}

      const group_width = ${t.params.groupSize.x}u;
      const group_height = ${t.params.groupSize.y}u;
      const group_depth = ${t.params.groupSize.z}u;

      @compute @workgroup_size(group_width, group_height, group_depth)
      fn main(
        ${params}
        ) {
        let o_group_index = ((${group_id}.z * ${num_groups}.y) + ${group_id}.y) * ${num_groups}.x + ${group_id}.x;
        let o_global_index = o_group_index * ${invocationsPerGroup}u + ${local_index};
        var o: Outputs;
        o.local_id = ${local_id};
        o.local_index = ${local_index};
        o.global_id = ${global_id};
        o.group_id = ${group_id};
        o.num_groups = ${num_groups};
        ${linear_indexing ? `o.global_index = ${global_index};` : ``}
        ${linear_indexing ? `o.group_index = ${group_index};` : ``}
        outputs[o_global_index] = o;
      }
    `;

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

    // Offsets are in u32 size units
    const kLocalIdOffset = 0;
    const kLocalIndexOffset = 3;
    const kGlobalIdOffset = 4;
    const kGroupIdOffset = 8;
    const kNumGroupsOffset = 12;
    const kGlobalIndexOffset = 15;
    const kGroupIndexOffset = 16;
    const kOutputElementSize = linear_indexing ? 20 : 16;

    // Create the output buffers.
    const outputBuffer = t.createBufferTracked({
      size: totalInvocations * kOutputElementSize * 4,
      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();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    switch (t.params.dispatch) {
      case 'direct':
        pass.dispatchWorkgroups(t.params.numGroups.x, t.params.numGroups.y, t.params.numGroups.z);
        break;
      case 'indirect': {
        const dispatchBuffer = t.createBufferTracked({
          size: 3 * Uint32Array.BYTES_PER_ELEMENT,
          usage: GPUBufferUsage.INDIRECT,
          mappedAtCreation: true,
        });
        const dispatchData = new Uint32Array(dispatchBuffer.getMappedRange());
        dispatchData[0] = t.params.numGroups.x;
        dispatchData[1] = t.params.numGroups.y;
        dispatchData[2] = t.params.numGroups.z;
        dispatchBuffer.unmap();
        pass.dispatchWorkgroupsIndirect(dispatchBuffer, 0);
        break;
      }
    }
    pass.end();
    t.queue.submit([encoder.finish()]);

    type vec3 = { x: number; y: number; z: number };

    // Helper to check that the vec3<u32> value at each index of the provided `output` buffer
    // matches the expected value for that invocation, as generated by the `getBuiltinValue`
    // function. The `name` parameter is the builtin name, used for error messages.
    const checkEachIndex = (output: Uint32Array) => {
      // Loop over workgroups.
      for (let gz = 0; gz < t.params.numGroups.z; gz++) {
        for (let gy = 0; gy < t.params.numGroups.y; gy++) {
          for (let gx = 0; gx < t.params.numGroups.x; gx++) {
            // Loop over invocations within a group.
            for (let lz = 0; lz < t.params.groupSize.z; lz++) {
              for (let ly = 0; ly < t.params.groupSize.y; ly++) {
                for (let lx = 0; lx < t.params.groupSize.x; lx++) {
                  const groupIndex = (gz * t.params.numGroups.y + gy) * t.params.numGroups.x + gx;
                  const localIndex = (lz * t.params.groupSize.y + ly) * t.params.groupSize.x + lx;
                  const globalIndex = groupIndex * invocationsPerGroup + localIndex;
                  const globalOffset = globalIndex * kOutputElementSize;
                  const gidX = gx * t.params.groupSize.x + lx;
                  const gidY = gy * t.params.groupSize.y + ly;
                  const gidZ = gz * t.params.groupSize.z + lz;
                  const globalLinearIndex =
                    gidX +
                    gidY * t.params.groupSize.x * t.params.numGroups.x +
                    gidZ *
                      t.params.groupSize.x *
                      t.params.numGroups.x *
                      t.params.groupSize.y *
                      t.params.numGroups.y;
                  const groupLinearIndex =
                    gx +
                    gy * t.params.numGroups.x +
                    gz * t.params.numGroups.x * t.params.numGroups.y;

                  const expectEqual = (name: string, expected: number, actual: number) => {
                    if (actual !== expected) {
                      return new Error(
                        `${name} failed at group(${gx},${gy},${gz}) local(${lx},${ly},${lz}))\n` +
                          `    expected: ${expected}\n` +
                          `    got:      ${actual}`
                      );
                    }
                    return undefined;
                  };

                  const checkVec3Value = (name: string, fieldOffset: number, expected: vec3) => {
                    const offset = globalOffset + fieldOffset;
                    return (
                      expectEqual(`${name}.x`, expected.x, output[offset + 0]) ||
                      expectEqual(`${name}.y`, expected.y, output[offset + 1]) ||
                      expectEqual(`${name}.z`, expected.z, output[offset + 2])
                    );
                  };

                  const error =
                    checkVec3Value('local_id', kLocalIdOffset, { x: lx, y: ly, z: lz }) ||
                    checkVec3Value('global_id', kGlobalIdOffset, { x: gidX, y: gidY, z: gidZ }) ||
                    checkVec3Value('group_id', kGroupIdOffset, { x: gx, y: gy, z: gz }) ||
                    checkVec3Value('num_groups', kNumGroupsOffset, t.params.numGroups) ||
                    expectEqual(
                      'local_index',
                      localIndex,
                      output[globalOffset + kLocalIndexOffset]
                    ) ||
                    (linear_indexing &&
                      expectEqual(
                        'global_index',
                        globalLinearIndex,
                        output[globalOffset + kGlobalIndexOffset]
                      )) ||
                    (linear_indexing &&
                      expectEqual(
                        'group_index',
                        groupLinearIndex,
                        output[globalOffset + kGroupIndexOffset]
                      ));
                  if (error) {
                    return error;
                  }
                }
              }
            }
          }
        }
      }
      return undefined;
    };

    t.expectGPUBufferValuesPassCheck(outputBuffer, outputData => checkEachIndex(outputData), {
      type: Uint32Array,
      typedLength: outputBuffer.size / 4,
    });
  });

/**
 * @returns The population count of input.
 *
 * @param input Treated as an unsigned 32-bit integer
 */
function popcount(input: number): number {
  let n = input;
  n = n - ((n >> 1) & 0x55555555);
  n = (n & 0x33333333) + ((n >> 2) & 0x33333333);
  return (((n + (n >> 4)) & 0xf0f0f0f) * 0x1010101) >> 24;
}

function ErrorMsg(msg: string, got: number, expected: number): string {
  return `${msg}:\n-      got: ${got}\n- expected: ${expected}`;
}

/**
 * Checks that the subgroup size and ballot buffers are consistent.
 *
 * This function checks that all invocations see a consistent value for
 * subgroup_size and that all ballots are less than or equal to that value.
 *
 * @param subgroupSizes The subgroup_size buffer
 * @param ballotSizes The ballot buffer
 * @param min The minimum subgroup size allowed
 * @param max The maximum subgroup size allowed
 * @param invocations The number of invocations in a workgroup
 */
function checkSubgroupSizeConsistency(
  subgroupSizes: Uint32Array,
  ballotSizes: Uint32Array,
  min: number,
  max: number,
  invocations: number
): Error | undefined {
  const subgroupSize = subgroupSizes[0];
  if (popcount(subgroupSize) !== 1) {
    return new Error(`Subgroup size '${subgroupSize}' is not a power of two`);
  }
  if (subgroupSize < min) {
    return new Error(`Subgroup size '${subgroupSize}' is less than minimum '${min}'`);
  }
  if (max < subgroupSize) {
    return new Error(`Subgroup size '${subgroupSize}' is greater than maximum '${max}'`);
  }

  // Check that remaining invocations record a consistent subgroup size.
  for (let i = 1; i < subgroupSizes.length; i++) {
    if (subgroupSizes[i] !== subgroupSize) {
      return new Error(
        ErrorMsg(`Invocation ${i}: subgroup size inconsistency`, subgroupSizes[i], subgroupSize)
      );
    }
  }

  for (let i = 0; i < ballotSizes.length; i++) {
    if (ballotSizes[i] > subgroupSize) {
      return new Error(
        `Invocation ${i}, subgroup size, ${ballotSizes[i]}, is greater than built-in value, ${subgroupSize}`
      );
    }
  }

  return undefined;
}

/**
 * Returns a WGSL function that generates linear local id
 *
 * Using (0,1,2) generates the standard local linear id.
 * Changing the order of p0, p1, and p2 changes the linearity.
 *
 * Assumes p0, p1, and p2 are not repeated and in range [0, 2].
 * @param p0 The index used for the x-dimension
 * @param p1 The index used for the y-dimension
 * @param p2 The index used for the z-dimension
 * @param sizes An array of workgroup sizes for each dimension.
 */
function genLID(p0: number, p1: number, p2: number, sizes: readonly number[]): string {
  return `
fn getLID(lid : vec3u) -> u32 {
  let p0 = lid[${p0}];
  let p1 = lid[${p1}] * ${sizes[p0]};
  let p2 = lid[${p2}] * ${sizes[p0]} * ${sizes[p1]};
  return p0 + p1 + p2;
}`;
}

const kWGSizes = [
  [1, 1, 1],
  [4, 1, 1],
  [8, 1, 1],
  [16, 1, 1],
  [32, 1, 1],
  [64, 1, 1],
  [128, 1, 1],
  [256, 1, 1],
  [1, 4, 1],
  [1, 8, 1],
  [1, 16, 1],
  [1, 32, 1],
  [1, 64, 1],
  [1, 128, 1],
  [1, 256, 1],
  [1, 1, 4],
  [1, 1, 8],
  [1, 1, 16],
  [1, 1, 32],
  [1, 1, 64],
  [3, 3, 3],
  [4, 4, 4],
  [16, 16, 1],
  [16, 1, 16],
  [1, 16, 16],
  [15, 3, 3],
  [3, 15, 3],
  [3, 3, 15],
  [17, 5, 2],
  [17, 4, 2],
  [15, 2, 8],
] as const;

g.test('subgroup_size')
  .desc('Tests subgroup_size values')
  .params(u =>
    u
      .combine('sizes', kWGSizes)
      .beginSubcases()
      .combine('numWGs', [1, 2] as const)
      .combine('lid', [
        [0, 1, 2],
        [0, 2, 1],
        [1, 0, 2],
        [1, 2, 0],
        [2, 0, 1],
        [2, 1, 0],
      ] as const)
  )
  .fn(async t => {
    t.skipIfDeviceDoesNotHaveFeature('subgroups' as GPUFeatureName);
    const subgroupMinSize = t.device.adapterInfo.subgroupMinSize!;
    const subgroupMaxSize = t.device.adapterInfo.subgroupMaxSize!;

    const wgx = t.params.sizes[0];
    const wgy = t.params.sizes[1];
    const wgz = t.params.sizes[2];
    const lid = t.params.lid;
    const wgThreads = wgx * wgy * wgz;

    // Compatibility mode has lower workgroup limits.
    const {
      maxComputeInvocationsPerWorkgroup,
      maxComputeWorkgroupSizeX,
      maxComputeWorkgroupSizeY,
      maxComputeWorkgroupSizeZ,
    } = t.device.limits;
    t.skipIf(
      maxComputeInvocationsPerWorkgroup < wgThreads ||
        maxComputeWorkgroupSizeX < t.params.sizes[0] ||
        maxComputeWorkgroupSizeY < t.params.sizes[1] ||
        maxComputeWorkgroupSizeZ < t.params.sizes[2],
      'Workgroup size too large'
    );

    const wgsl = `
enable subgroups;

const stride = ${wgThreads};

${genLID(lid[0], lid[1], lid[2], t.params.sizes)}

@group(0) @binding(0)
var<storage, read_write> output : array<u32>;

@group(0) @binding(1)
var<storage, read_write> compare : array<u32>;

@compute @workgroup_size(${wgx}, ${wgy}, ${wgz})
fn main(@builtin(subgroup_size) size : u32,
        @builtin(workgroup_id) wgid : vec3u,
        @builtin(local_invocation_id) local_id : vec3u) {
  // Remap local ids according to test linearity.
  let lid = getLID(local_id);

  output[lid + wgid.x * stride] = size;
  let ballot = countOneBits(subgroupBallot(true));
  let ballotSize = ballot[0] + ballot[1] + ballot[2] + ballot[3];
  compare[lid + wgid.x * stride] = ballotSize;
}`;

    const numInvocations = wgThreads * t.params.numWGs;
    const sizesBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numInvocations, x => 0)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(sizesBuffer);
    const compareBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numInvocations, x => 0)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(compareBuffer);

    const pipeline = t.device.createComputePipeline({
      layout: 'auto',
      compute: {
        module: t.device.createShaderModule({
          code: wgsl,
        }),
        entryPoint: 'main',
      },
    });
    const bg = t.device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        {
          binding: 0,
          resource: {
            buffer: sizesBuffer,
          },
        },
        {
          binding: 1,
          resource: {
            buffer: compareBuffer,
          },
        },
      ],
    });

    const encoder = t.device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bg);
    pass.dispatchWorkgroups(t.params.numWGs, 1, 1);
    pass.end();
    t.queue.submit([encoder.finish()]);

    const sizesReadback = await t.readGPUBufferRangeTyped(sizesBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numInvocations,
      method: 'copy',
    });
    const sizesData: Uint32Array = sizesReadback.data;

    const compareReadback = await t.readGPUBufferRangeTyped(compareBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numInvocations,
      method: 'copy',
    });
    const compareData: Uint32Array = compareReadback.data;

    t.expectOK(
      checkSubgroupSizeConsistency(
        sizesData,
        compareData,
        subgroupMinSize,
        subgroupMaxSize,
        wgThreads
      )
    );
  });

/**
 * Checks the consistency of subgroup_invocation_id builtin values
 *
 * Creates a ballot out consisting of all invocations sharing the same generated
 * subgroup id. Checks that the ballot contains at most subgroupSize invocations
 * and that the invocations are tightly packed from the lowest id.
 * @param data The subgroup_invocation_id output data
 * @param ids The representative subgroup ids for each invocation
 * @param subgroupSize The subgroup size
 * @param invocations The number of invocations per workgroup
 * @param numWGs The number of workgroups
 */
function checkSubgroupInvocationIdConsistency(
  data: Uint32Array,
  ids: Uint32Array,
  subgroupSize: number,
  invocations: number,
  numWGs: number
): Error | undefined {
  for (let wg = 0; wg < numWGs; wg++) {
    // Tracks an effective ballot of each subgroup based on the representative id
    // (global id of invocation 0 in the subgroup).
    const mappings = new Map<number, bigint>();
    for (let i = 0; i < invocations; i++) {
      const idx = i + invocations * wg;
      const subgroup_id = ids[idx];
      if (subgroup_id === 999) {
        return new Error(`Invocation ${i}: no data`);
      }
      let v = mappings.get(subgroup_id) ?? 0n;
      v |= 1n << BigInt(data[idx]);
      mappings.set(subgroup_id, v);
    }

    for (const value of mappings.entries()) {
      const id = value[0];
      const ballot = value[1];

      let onebits = popcount(Number(BigInt.asUintN(32, ballot)));
      onebits += popcount(Number(BigInt.asUintN(32, ballot >> 32n)));
      onebits += popcount(Number(BigInt.asUintN(32, ballot >> 64n)));
      onebits += popcount(Number(BigInt.asUintN(32, ballot >> 96n)));
      if (onebits > subgroupSize) {
        return new Error(
          `Subgroup including invocation ${id} is too large, ${onebits}, for subgroup size, ${subgroupSize}`
        );
      }

      const ballotP1 = ballot + 1n;
      if ((ballot & ballotP1) !== 0n) {
        return new Error(
          `Subgroup including invocation ${id} has non-continguous ids: ${ballot.toString(2)}`
        );
      }
    }
  }

  return undefined;
}

g.test('subgroup_invocation_id')
  .desc(
    'Tests subgroup_invocation_id values. No mapping between local_invocation_index and subgroup_invocation_id can be relied upon.'
  )
  .params(u =>
    u
      .combine('sizes', kWGSizes)
      .beginSubcases()
      .combine('numWGs', [1, 2] as const)
      .combine('lid', [
        [0, 1, 2],
        [0, 2, 1],
        [1, 0, 2],
        [1, 2, 0],
        [2, 0, 1],
        [2, 1, 0],
      ] as const)
  )
  .fn(async t => {
    t.skipIfDeviceDoesNotHaveFeature('subgroups' as GPUFeatureName);
    const wgx = t.params.sizes[0];
    const wgy = t.params.sizes[1];
    const wgz = t.params.sizes[2];
    const lid = t.params.lid;
    const wgThreads = wgx * wgy * wgz;

    // Compatibility mode has lower workgroup limits.
    const {
      maxComputeInvocationsPerWorkgroup,
      maxComputeWorkgroupSizeX,
      maxComputeWorkgroupSizeY,
      maxComputeWorkgroupSizeZ,
    } = t.device.limits;
    t.skipIf(
      maxComputeInvocationsPerWorkgroup < wgThreads ||
        maxComputeWorkgroupSizeX < wgx ||
        maxComputeWorkgroupSizeY < wgy ||
        maxComputeWorkgroupSizeZ < wgz,
      'Workgroup size too large'
    );

    const wgsl = `
enable subgroups;

const stride = ${wgThreads};

${genLID(lid[0], lid[1], lid[2], t.params.sizes)}

@group(0) @binding(0)
var<storage, read_write> output : array<u32>;

// This var stores the global id of invocation 0 in the subgroup.
@group(0) @binding(1)
var<storage, read_write> subgroup_ids : array<u32>;

@group(0) @binding(2)
var<storage, read_write> sizes : array<u32>;

@compute @workgroup_size(${wgx}, ${wgy}, ${wgz})
fn main(@builtin(subgroup_size) size : u32,
        @builtin(subgroup_invocation_id) id : u32,
        @builtin(workgroup_id) wgid : vec3u,
        @builtin(local_invocation_id) local_id : vec3u) {
  // Remap local ids according to test linearity.
  let lid = getLID(local_id);

  // Representative subgroup_id value.
  let gid = lid + stride * wgid.x;

  let b = subgroupBroadcast(gid, 0);
  output[gid] = id;
  subgroup_ids[gid] = b;
  if (lid == 0) {
    sizes[wgid.x] = size;
  }
}`;

    const numInvocations = wgThreads * t.params.numWGs;
    const outputBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numInvocations, x => 0)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(outputBuffer);
    const placeholderValue = 999;
    const idsBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numInvocations, x => placeholderValue)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(idsBuffer);
    const sizeBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(t.params.numWGs, x => 0)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(sizeBuffer);

    const pipeline = t.device.createComputePipeline({
      layout: 'auto',
      compute: {
        module: t.device.createShaderModule({
          code: wgsl,
        }),
        entryPoint: 'main',
      },
    });
    const bg = t.device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        {
          binding: 0,
          resource: {
            buffer: outputBuffer,
          },
        },
        {
          binding: 1,
          resource: {
            buffer: idsBuffer,
          },
        },
        {
          binding: 2,
          resource: {
            buffer: sizeBuffer,
          },
        },
      ],
    });

    const encoder = t.device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bg);
    pass.dispatchWorkgroups(t.params.numWGs, 1, 1);
    pass.end();
    t.queue.submit([encoder.finish()]);

    const sizeReadback = await t.readGPUBufferRangeTyped(sizeBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: 1,
      method: 'copy',
    });
    const sizeData: Uint32Array = sizeReadback.data;

    const outputReadback = await t.readGPUBufferRangeTyped(outputBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numInvocations,
      method: 'copy',
    });
    const outputData: Uint32Array = outputReadback.data;

    const idsReadback = await t.readGPUBufferRangeTyped(idsBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numInvocations,
      method: 'copy',
    });
    const idsData: Uint32Array = idsReadback.data;

    t.expectOK(
      checkSubgroupInvocationIdConsistency(
        outputData,
        idsData,
        sizeData[0],
        wgThreads,
        t.params.numWGs
      )
    );
  });

const skipValue = 0xffff0000;

/**
 * Checks subgroup_id consistency
 *
 * @param outputData An array of vec4u
 *                   * 0: comparison of subgroup_id among subgroup
 *                   * 1: comparison of subgroup_id < num_subgroups
 *                   * 2: subgroup_id (for first member) or skipValue
 *                   * 3: unused
 * @param wgSize Invocations in the workgroup
 * @param numWGs Number of workgroups
 */
function checkSubgroupIdConsistency(
  outputData: Uint32Array,
  wgSize: number,
  numWGs: number
): Error | undefined {
  for (let wg = 0; wg < numWGs; wg++) {
    // Max wgSize is 256 and min subgroup size is 4
    const seen = new Array(Math.ceil(wgSize / 4));
    seen.fill(0);
    for (let inv = 0; inv < wgSize; inv++) {
      const gid = wg * wgSize + inv;
      const outputIdx = gid * 4;
      const compare = outputData[outputIdx];
      const in_range = outputData[outputIdx + 1];
      const sid = outputData[outputIdx + 2];

      if (compare !== 1) {
        return new Error(
          `Invocation ${gid}: not all invocations in subgroup have same subgroup_id: ${compare}`
        );
      }
      if (in_range !== 1) {
        return new Error(
          `Invocation ${gid}: subgroup_id out of range of num_subgroups: ${in_range}`
        );
      }

      if (sid !== skipValue) {
        if (seen[sid] !== 0) {
          return new Error(`Invocation ${gid}: subgroup_id reused among different subgroups`);
        }
        seen[sid] = 1;
      }
    }

    const firstZero = seen.findIndex(ele => ele === 0);
    const lastOne = seen.findLastIndex(ele => ele === 1);
    if (firstZero !== -1 && firstZero < lastOne) {
      return new Error(`Subgroup id values are not densely packed: missing ${firstZero}`);
    }
  }

  return undefined;
}

g.test('subgroup_id')
  .desc(
    'Tests subgroup_id values. No mapping between local_invocation_index and subgroup_id can be relied upon.'
  )
  .params(u =>
    u
      .combine('sizes', kWGSizes)
      .beginSubcases()
      .combine('numWGs', [1, 2] as const)
      .combine('lid', [
        [0, 1, 2],
        [0, 2, 1],
        [1, 0, 2],
        [1, 2, 0],
        [2, 0, 1],
        [2, 1, 0],
      ] as const)
  )
  .fn(async t => {
    t.skipIfDeviceDoesNotHaveFeature('subgroups' as GPUFeatureName);
    t.skipIfLanguageFeatureNotSupported('subgroup_id');
    const wgx = t.params.sizes[0];
    const wgy = t.params.sizes[1];
    const wgz = t.params.sizes[2];
    const lid = t.params.lid;
    const wgThreads = wgx * wgy * wgz;

    // Compatibility mode has lower workgroup limits.
    const {
      maxComputeInvocationsPerWorkgroup,
      maxComputeWorkgroupSizeX,
      maxComputeWorkgroupSizeY,
      maxComputeWorkgroupSizeZ,
    } = t.device.limits;
    t.skipIf(
      maxComputeInvocationsPerWorkgroup < wgThreads ||
        maxComputeWorkgroupSizeX < wgx ||
        maxComputeWorkgroupSizeY < wgy ||
        maxComputeWorkgroupSizeZ < wgz,
      'Workgroup size too large'
    );

    const wgsl = `
enable subgroups;
requires subgroup_id;

const stride = ${wgThreads};

${genLID(lid[0], lid[1], lid[2], t.params.sizes)}

@group(0) @binding(0)
var<storage, read_write> output : array<vec4u>;

@compute @workgroup_size(${wgx}, ${wgy}, ${wgz})
fn main(@builtin(local_invocation_id) local_id : vec3u,
        @builtin(workgroup_id) wgid : vec3u,
        @builtin(subgroup_id) sid : u32,
        @builtin(num_subgroups) num_subgroups : u32) {
  // Remapped local id.
  let lid = getLID(local_id);

  let gid = lid + stride * wgid.x;

  // Is the subgroup_id equivalent for all members?
  let broadcast_id = subgroupBroadcastFirst(sid);
  let compare = subgroupAll(broadcast_id == sid);

  // Is subgroup_id in the range of num_subgroups?
  let in_range = sid < num_subgroups;

  var out_sid = ${skipValue}u;
  if subgroupElect() {
    out_sid = sid;
  }

  output[gid] = vec4u(
    select(0u, 1u, compare),
    select(0u, 1u, in_range),
    out_sid,
    0);
}
`;

    const numInvocations = wgThreads * t.params.numWGs;
    const numUints = 4 * numInvocations;
    const placeholderValue = 999;
    const outputBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numUints, x => placeholderValue)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(outputBuffer);

    const pipeline = t.device.createComputePipeline({
      layout: 'auto',
      compute: {
        module: t.device.createShaderModule({
          code: wgsl,
        }),
        entryPoint: 'main',
      },
    });
    const bg = t.device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        {
          binding: 0,
          resource: {
            buffer: outputBuffer,
          },
        },
      ],
    });

    const encoder = t.device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bg);
    pass.dispatchWorkgroups(t.params.numWGs, 1, 1);
    pass.end();
    t.queue.submit([encoder.finish()]);

    const outputReadback = await t.readGPUBufferRangeTyped(outputBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numUints,
      method: 'copy',
    });
    const outputData: Uint32Array = outputReadback.data;

    t.expectOK(checkSubgroupIdConsistency(outputData, wgThreads, t.params.numWGs));
  });

/**
 * Checks num_subgroups consistency
 *
 * @param countData An array with numWGs elements containing the counted number of subgroups
 * @param outputData An array numWGs * wgSize elements containing the value of num_subgroups
 * @param wgSize Number of invocations in the workgroup
 * @param numWGs Number of workgroups
 */
function checkNumSubgroupsConsistency(
  countData: Uint32Array,
  outputData: Uint32Array,
  wgSize: number,
  numWGs: number
): Error | undefined {
  for (let wg = 0; wg < numWGs; wg++) {
    const count = countData[wg];
    const slice = outputData.slice(wg * wgSize, (wg + 1) * wgSize);
    const index = slice.findIndex(ele => ele !== count);
    if (index !== -1) {
      return new Error(`Workgroup ${wg}: inconsistent num_subgroups:
- expected: ${count}
-      got: ${slice[index]}`);
    }
  }

  return undefined;
}

g.test('num_subgroups')
  .desc('Tests num_subgroups values.')
  .params(u =>
    u
      .combine('sizes', kWGSizes)
      .beginSubcases()
      .combine('numWGs', [1, 2] as const)
      .combine('lid', [
        [0, 1, 2],
        [0, 2, 1],
        [1, 0, 2],
        [1, 2, 0],
        [2, 0, 1],
        [2, 1, 0],
      ] as const)
  )
  .fn(async t => {
    t.skipIfDeviceDoesNotHaveFeature('subgroups' as GPUFeatureName);
    t.skipIfLanguageFeatureNotSupported('subgroup_id');
    const wgx = t.params.sizes[0];
    const wgy = t.params.sizes[1];
    const wgz = t.params.sizes[2];
    const lid = t.params.lid;
    const wgThreads = wgx * wgy * wgz;

    // Compatibility mode has lower workgroup limits.
    const {
      maxComputeInvocationsPerWorkgroup,
      maxComputeWorkgroupSizeX,
      maxComputeWorkgroupSizeY,
      maxComputeWorkgroupSizeZ,
    } = t.device.limits;
    t.skipIf(
      maxComputeInvocationsPerWorkgroup < wgThreads ||
        maxComputeWorkgroupSizeX < wgx ||
        maxComputeWorkgroupSizeY < wgy ||
        maxComputeWorkgroupSizeZ < wgz,
      'Workgroup size too large'
    );

    const wgsl = `
enable subgroups;
requires subgroup_id;

const stride = ${wgThreads};

${genLID(lid[0], lid[1], lid[2], t.params.sizes)}

@group(0) @binding(0)
var<storage, read_write> numSubgroups : array<u32>;

@group(0) @binding(1)
var<storage, read_write> output : array<u32>;

var<workgroup> count : atomic<u32>;

@compute @workgroup_size(${wgx}, ${wgy}, ${wgz})
fn main(@builtin(local_invocation_id) local_id : vec3u,
        @builtin(workgroup_id) wgid : vec3u,
        @builtin(subgroup_id) sid : u32,
        @builtin(num_subgroups) num_subgroups : u32) {
  // Remapped local id.
  let lid = getLID(local_id);

  let gid = lid + stride * wgid.x;

  if subgroupElect() {
    atomicAdd(&count, 1);
  }

  workgroupBarrier();

  if lid == 0 {
    numSubgroups[wgid.x] = atomicLoad(&count);
  }

  output[gid] = num_subgroups;
}
`;

    const numInvocations = wgThreads * t.params.numWGs;
    const placeholderValue = 999;
    const countBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(t.params.numWGs, x => placeholderValue)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(countBuffer);
    const outputBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numInvocations * 4, x => placeholderValue)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );
    t.trackForCleanup(outputBuffer);

    const pipeline = t.device.createComputePipeline({
      layout: 'auto',
      compute: {
        module: t.device.createShaderModule({
          code: wgsl,
        }),
        entryPoint: 'main',
      },
    });
    const bg = t.device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        {
          binding: 0,
          resource: {
            buffer: countBuffer,
          },
        },
        {
          binding: 1,
          resource: {
            buffer: outputBuffer,
          },
        },
      ],
    });

    const encoder = t.device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bg);
    pass.dispatchWorkgroups(t.params.numWGs, 1, 1);
    pass.end();
    t.queue.submit([encoder.finish()]);

    const countReadback = await t.readGPUBufferRangeTyped(countBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: t.params.numWGs,
      method: 'copy',
    });
    const countData: Uint32Array = countReadback.data;
    const outputReadback = await t.readGPUBufferRangeTyped(outputBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numInvocations,
      method: 'copy',
    });
    const outputData: Uint32Array = outputReadback.data;

    t.expectOK(checkNumSubgroupsConsistency(countData, outputData, wgThreads, t.params.numWGs));
  });

g.test('subgroup_size_attribute')
  .desc(
    `Tests that all power-of-two values in [subgroupMinSize, subgroupMaxSize] can be used as
    the @subgroup_size attribute in a simple compute pipeline. The value of the subgroup_size
    builtin must equal the value of the @subgroup_size attribute.`
  )
  .params(u =>
    u
      .combine('subgroupSize', [4, 8, 16, 32, 64, 128] as const)
      .combine('numWorkGroups', [1, 2] as const)
      .combine('numSubgroups', [1, 2, 4] as const)
  )
  .fn(async t => {
    t.skipIfDeviceDoesNotHaveFeature('subgroup-size-control');

    const { subgroupSize, numWorkGroups, numSubgroups } = t.params;

    const subgroupMinSize = t.device.adapterInfo.subgroupMinSize!;
    const subgroupMaxSize = t.device.adapterInfo.subgroupMaxSize!;

    t.skipIf(
      subgroupSize < subgroupMinSize || subgroupSize > subgroupMaxSize,
      `subgroupSize ${subgroupSize} is outside device range [${subgroupMinSize}, ${subgroupMaxSize}]`
    );

    const wgx = subgroupSize * numSubgroups;
    t.skipIf(
      wgx > t.device.limits.maxComputeWorkgroupSizeX,
      `workgroup size ${wgx} exceeds maxComputeWorkgroupSizeX`
    );

    const wgsl = `
enable subgroups;
enable subgroup_size_control;

@group(0) @binding(0)
var<storage, read_write> output : array<u32>;

@compute @workgroup_size(${wgx}, 1, 1) @subgroup_size(${subgroupSize})
fn main(@builtin(subgroup_size) builtin_size : u32,
        @builtin(local_invocation_index) lid : u32,
        @builtin(workgroup_id) wgid : vec3u) {
  let gid = lid + wgid.x * ${wgx}u;
  // Store 1 if builtin subgroup_size matches the @subgroup_size attribute, 0 otherwise.
  output[gid] = select(0u, 1u, builtin_size == ${subgroupSize}u);
}`;

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

    const numInvocations = wgx * numWorkGroups;
    const outputBuffer = t.makeBufferWithContents(
      new Uint32Array([...iterRange(numInvocations, x => 0)]),
      GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
    );

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

    const encoder = t.device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bg);
    pass.dispatchWorkgroups(numWorkGroups, 1, 1);
    pass.end();
    t.queue.submit([encoder.finish()]);

    const outputReadback = await t.readGPUBufferRangeTyped(outputBuffer, {
      srcByteOffset: 0,
      type: Uint32Array,
      typedLength: numInvocations,
      method: 'copy',
    });
    const outputData: Uint32Array = outputReadback.data;

    for (let i = 0; i < numInvocations; i++) {
      if (outputData[i] !== 1) {
        t.fail(
          `@subgroup_size(${subgroupSize}): invocation ${i} has builtin subgroup_size != ${subgroupSize}`
        );
        break;
      }
    }
  });
