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);
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;
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;
}
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',
},
});
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;
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 } }],
});
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 };
const checkEachIndex = (output: Uint32Array) => {
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++) {
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}'`);
}
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;
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++) {
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;
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++) {
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;
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;
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;
}
}
});
