Skip to content

Commit

Permalink
add f16 to webgpu/shader/types.ts utilities
Browse files Browse the repository at this point in the history 
Udpate affected tests;
- webgpu:shader,execution,zero_init:
  skip f16 cases
- webgpu:shader,execution,robust_access:
  generalize to cover most f16 cases

Bug: #3405
  • Loading branch information
@dneto0
dneto0 committed Feb 15, 2024
1 parent 7ad3a97 commit 704bde6
Show file tree
Hide file tree
Showing 4 changed files with 118 additions and 33 deletions.
29 changes: 24 additions & 5 deletions src/webgpu/shader/execution/robust_access.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ TODO: add tests to check that textureLoad operations stay in-bounds.

import { makeTestGroup } from '../../../common/framework/test_group.js';
import { assert } from '../../../common/util/util.js';
import { Float16Array } from '../../../external/petamoriken/float16/float16.js';
import { GPUTest } from '../../gpu_test.js';
import { align } from '../../util/math.js';
import { generateTypes, supportedScalarTypes, supportsAtomics } from '../types.js';
Expand All @@ -25,6 +26,7 @@ const kMinI32 = -0x8000_0000;
*/
async function runShaderTest(
t: GPUTest,
enables: string,
stage: GPUShaderStageFlags,
testSource: string,
layout: GPUPipelineLayout,
Expand All @@ -41,7 +43,7 @@ async function runShaderTest(
usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE,
});

const source = `
const source = `${enables}
struct Constants {
zero: u32
};
Expand Down Expand Up @@ -96,10 +98,12 @@ fn main() {
/** Fill an ArrayBuffer with sentinel values, except clear a region to zero. */
function testFillArrayBuffer(
array: ArrayBuffer,
type: 'u32' | 'i32' | 'f32',
type: 'u32' | 'i32' | 'f16' | 'f32',
{ zeroByteStart, zeroByteCount }: { zeroByteStart: number; zeroByteCount: number }
) {
const constructor = { u32: Uint32Array, i32: Int32Array, f32: Float32Array }[type];
const constructor = { u32: Uint32Array, i32: Int32Array, f16: Float16Array, f32: Float32Array }[
type
];
assert(zeroByteCount % constructor.BYTES_PER_ELEMENT === 0);
new constructor(array).fill(42);
new constructor(array, zeroByteStart, zeroByteCount / constructor.BYTES_PER_ELEMENT).fill(0);
Expand Down Expand Up @@ -168,10 +172,15 @@ g.test('linear_memory')
{ shadowingMode: 'function-scope' },
])
.expand('isAtomic', p => (supportsAtomics(p) ? [false, true] : [false]))
.beginSubcases()
.expand('baseType', supportedScalarTypes)
.beginSubcases()
.expandWithParams(generateTypes)
)
.beforeAllSubcases(t => {
if (t.params.baseType === 'f16') {
t.selectDeviceOrSkipTestCase('shader-f16');
}
})
.fn(async t => {
const {
addressSpace,
Expand All @@ -189,6 +198,13 @@ g.test('linear_memory')
assert(_kTypeInfo !== undefined, 'not an indexable type');
assert('arrayLength' in _kTypeInfo);

if (baseType === 'f16' && addressSpace === 'uniform' && containerType === 'array') {
// Array elements must be aligned to 16 bytes, but the logic in generateTypes
// creates an array of vec4 of the baseType. But for f16 that's only 8 bytes.
// We would need to write more complex logic for that.
t.skip('TODO: Test logic does not handle array of f16 in the uniform address space');
}

let usesCanary = false;
let globalSource = '';
let testFunctionSource = '';
Expand Down Expand Up @@ -429,6 +445,8 @@ fn runTest() -> u32 {
],
});

const enables = t.params.baseType === 'f16' ? 'enable f16;' : '';

// Run it.
if (bufferBindingSize !== undefined && baseType !== 'bool') {
const expectedData = new ArrayBuffer(testBufferSize);
Expand All @@ -450,6 +468,7 @@ fn runTest() -> u32 {
// Run the shader, accessing the buffer.
await runShaderTest(
t,
enables,
GPUShaderStage.COMPUTE,
testSource,
layout,
Expand All @@ -475,6 +494,6 @@ fn runTest() -> u32 {
bufferBindingEnd
);
} else {
await runShaderTest(t, GPUShaderStage.COMPUTE, testSource, layout, []);
await runShaderTest(t, enables, GPUShaderStage.COMPUTE, testSource, layout, []);
}
});
4 changes: 4 additions & 0 deletions src/webgpu/shader/execution/zero_init.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,10 @@ g.test('compute,zero_init')
? [true, false]
: [false]) {
for (const scalarType of supportedScalarTypes({ isAtomic, ...p })) {
// Fewer subcases: supportedScalarTypes was expanded to include f16
// but that may take too much time. It would require more complex code.
if (scalarType === 'f16') continue;

// Fewer subcases: For nested types, skip atomic u32 and non-atomic i32.
if (p._containerDepth > 0) {
if (scalarType === 'u32' && isAtomic) continue;
Expand Down
113 changes: 86 additions & 27 deletions src/webgpu/shader/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,20 @@ import { keysOf } from '../../common/util/data_tables.js';
import { assert } from '../../common/util/util.js';
import { align } from '../util/math.js';

const kArrayLength = 3;
const kDefaultArrayLength = 3;

export type Requirement = 'never' | 'may' | 'must'; // never is the same as "must not"
export type ContainerType = 'scalar' | 'vector' | 'matrix' | 'atomic' | 'array';
export type ScalarType = 'i32' | 'u32' | 'f32' | 'bool';
export type ScalarType = 'i32' | 'u32' | 'f16' | 'f32' | 'bool';

export const HostSharableTypes = ['i32', 'u32', 'f32'] as const;
export const HostSharableTypes = ['i32', 'u32', 'f16', 'f32'] as const;

/** Info for each plain scalar type. */
export const kScalarTypeInfo =
/* prettier-ignore */ {
'i32': { layout: { alignment: 4, size: 4 }, supportsAtomics: true, arrayLength: 1, innerLength: 0 },
'u32': { layout: { alignment: 4, size: 4 }, supportsAtomics: true, arrayLength: 1, innerLength: 0 },
'f16': { layout: { alignment: 2, size: 2 }, supportsAtomics: false, arrayLength: 1, innerLength: 0, feature: 'shader-f16' },
'f32': { layout: { alignment: 4, size: 4 }, supportsAtomics: false, arrayLength: 1, innerLength: 0 },
'bool': { layout: undefined, supportsAtomics: false, arrayLength: 1, innerLength: 0 },
} as const;
Expand All @@ -24,29 +25,71 @@ export const kScalarTypes = keysOf(kScalarTypeInfo);
/** Info for each vecN<> container type. */
export const kVectorContainerTypeInfo =
/* prettier-ignore */ {
'vec2': { layout: { alignment: 8, size: 8 }, arrayLength: 2 , innerLength: 0 },
'vec3': { layout: { alignment: 16, size: 12 }, arrayLength: 3 , innerLength: 0 },
'vec4': { layout: { alignment: 16, size: 16 }, arrayLength: 4 , innerLength: 0 },
'vec2': { arrayLength: 2 , innerLength: 0 },
'vec3': { arrayLength: 3 , innerLength: 0 },
'vec4': { arrayLength: 4 , innerLength: 0 },
} as const;
/** List of all vecN<> container types. */
export const kVectorContainerTypes = keysOf(kVectorContainerTypeInfo);

/** Returns the vector layout for a given vector container and base type, or undefined if that base type has no layout */
function vectorLayout(
vectorContainer: 'vec2' | 'vec3' | 'vec4',
baseType: ScalarType
): { alignment: number; size: number } | undefined {
const n = kVectorContainerTypeInfo[vectorContainer].arrayLength;
const scalarLayout = kScalarTypeInfo[baseType].layout;
if (scalarLayout === undefined) {
return undefined;
}
if (n === 3) {
return { alignment: scalarLayout.alignment * 4, size: scalarLayout.size * 3 };
}
return { alignment: scalarLayout.alignment * n, size: scalarLayout.size * n };
}

/** Info for each matNxN<> container type. */
export const kMatrixContainerTypeInfo =
/* prettier-ignore */ {
'mat2x2': { layout: { alignment: 8, size: 16 }, arrayLength: 2, innerLength: 2 },
'mat3x2': { layout: { alignment: 8, size: 24 }, arrayLength: 3, innerLength: 2 },
'mat4x2': { layout: { alignment: 8, size: 32 }, arrayLength: 4, innerLength: 2 },
'mat2x3': { layout: { alignment: 16, size: 32 }, arrayLength: 2, innerLength: 3 },
'mat3x3': { layout: { alignment: 16, size: 48 }, arrayLength: 3, innerLength: 3 },
'mat4x3': { layout: { alignment: 16, size: 64 }, arrayLength: 4, innerLength: 3 },
'mat2x4': { layout: { alignment: 16, size: 32 }, arrayLength: 2, innerLength: 4 },
'mat3x4': { layout: { alignment: 16, size: 48 }, arrayLength: 3, innerLength: 4 },
'mat4x4': { layout: { alignment: 16, size: 64 }, arrayLength: 4, innerLength: 4 },
'mat2x2': { arrayLength: 2, innerLength: 2 },
'mat3x2': { arrayLength: 3, innerLength: 2 },
'mat4x2': { arrayLength: 4, innerLength: 2 },
'mat2x3': { arrayLength: 2, innerLength: 3 },
'mat3x3': { arrayLength: 3, innerLength: 3 },
'mat4x3': { arrayLength: 4, innerLength: 3 },
'mat2x4': { arrayLength: 2, innerLength: 4 },
'mat3x4': { arrayLength: 3, innerLength: 4 },
'mat4x4': { arrayLength: 4, innerLength: 4 },
} as const;
/** List of all matNxN<> container types. */
export const kMatrixContainerTypes = keysOf(kMatrixContainerTypeInfo);

export const kMatrixContainerTypeLayoutInfo =
/* prettier-ignore */ {
'f16': {
'mat2x2': { layout: { alignment: 4, size: 8 } },
'mat3x2': { layout: { alignment: 4, size: 12 } },
'mat4x2': { layout: { alignment: 4, size: 16 } },
'mat2x3': { layout: { alignment: 8, size: 16 } },
'mat3x3': { layout: { alignment: 8, size: 24 } },
'mat4x3': { layout: { alignment: 8, size: 32 } },
'mat2x4': { layout: { alignment: 8, size: 16 } },
'mat3x4': { layout: { alignment: 8, size: 24 } },
'mat4x4': { layout: { alignment: 8, size: 32 } },
},
'f32': {
'mat2x2': { layout: { alignment: 8, size: 16 } },
'mat3x2': { layout: { alignment: 8, size: 24 } },
'mat4x2': { layout: { alignment: 8, size: 32 } },
'mat2x3': { layout: { alignment: 16, size: 32 } },
'mat3x3': { layout: { alignment: 16, size: 48 } },
'mat4x3': { layout: { alignment: 16, size: 64 } },
'mat2x4': { layout: { alignment: 16, size: 32 } },
'mat3x4': { layout: { alignment: 16, size: 48 } },
'mat4x4': { layout: { alignment: 16, size: 64 } },
}
} as const;

export type AddressSpace = 'storage' | 'uniform' | 'private' | 'function' | 'workgroup' | 'handle';
export type AccessMode = 'read' | 'write' | 'read_write';
export type Scope = 'module' | 'function';
Expand Down Expand Up @@ -189,21 +232,27 @@ export function* generateTypes({
for (const vectorType of kVectorContainerTypes) {
yield {
type: `${vectorType}<${scalarType}>`,
_kTypeInfo: { elementBaseType: baseType, ...kVectorContainerTypeInfo[vectorType] },
_kTypeInfo: {
elementBaseType: baseType,
...kVectorContainerTypeInfo[vectorType],
layout: vectorLayout(vectorType, scalarType as ScalarType),
},
};
}
}

if (containerType === 'matrix') {
// Matrices can only be f32.
if (baseType === 'f32') {
// Matrices can only be f16 or f32.
if (baseType === 'f16' || baseType === 'f32') {
for (const matrixType of kMatrixContainerTypes) {
const matrixInfo = kMatrixContainerTypeInfo[matrixType];
const matrixDimInfo = kMatrixContainerTypeInfo[matrixType];
const matrixLayoutInfo = kMatrixContainerTypeLayoutInfo[baseType][matrixType];
yield {
type: `${matrixType}<${scalarType}>`,
_kTypeInfo: {
elementBaseType: `vec${matrixInfo.innerLength}<${scalarType}>`,
...matrixInfo,
elementBaseType: `vec${matrixDimInfo.innerLength}<${scalarType}>`,
...matrixDimInfo,
...matrixLayoutInfo,
},
};
}
Expand All @@ -212,33 +261,43 @@ export function* generateTypes({

// Array types
if (containerType === 'array') {
// Buffer affective binding size must be a multiple of 4. Adjust array length as needed.
let arrayLength = kDefaultArrayLength;
if (
addressSpace === 'storage' &&
scalarInfo.layout !== undefined &&
scalarInfo.layout.alignment % 4 !== 0
) {
arrayLength = align(arrayLength, 4);
}

const arrayTypeInfo = {
elementBaseType: `${baseType}`,
arrayLength: kArrayLength,
arrayLength,
layout: scalarInfo.layout
? {
alignment: scalarInfo.layout.alignment,
size:
addressSpace === 'uniform'
? // Uniform storage class must have array elements aligned to 16.
kArrayLength *
arrayLength *
arrayStride({
...scalarInfo.layout,
alignment: 16,
})
: kArrayLength * arrayStride(scalarInfo.layout),
: arrayLength * arrayStride(scalarInfo.layout),
}
: undefined,
};

// Sized
if (addressSpace === 'uniform') {
yield {
type: `array<vec4<${scalarType}>,${kArrayLength}>`,
type: `array<vec4<${scalarType}>,${arrayLength}>`,
_kTypeInfo: arrayTypeInfo,
};
} else {
yield { type: `array<${scalarType},${kArrayLength}>`, _kTypeInfo: arrayTypeInfo };
yield { type: `array<${scalarType},${arrayLength}>`, _kTypeInfo: arrayTypeInfo };
}
// Unsized
if (addressSpace === 'storage') {
Expand Down Expand Up @@ -272,7 +331,7 @@ export function supportsAtomics(p: {
);
}

/** Generates an iterator of supported base types (i32/u32/f32/bool) */
/** Generates an iterator of supported base types (i32/u32/f16/f32/bool) */
export function* supportedScalarTypes(p: { isAtomic: boolean; addressSpace: string }) {
for (const scalarType of kScalarTypes) {
const info = kScalarTypeInfo[scalarType];
Expand Down
5 changes: 4 additions & 1 deletion src/webgpu/util/check_contents.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,10 @@ export function checkElementsEqual(
expected: TypedArrayBufferView
): ErrorWithExtra | undefined {
assert(actual.constructor === expected.constructor, 'TypedArray type mismatch');
assert(actual.length === expected.length, 'size mismatch');
assert(
actual.length === expected.length,
`length mismatch: expected ${expected.length} got ${actual.length}`
);

let failedElementsFirstMaybe: number | undefined = undefined;
/** Sparse array with `true` for elements that failed. */
Expand Down

0 comments on commit 704bde6

Please sign in to comment.