saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/webgpu] add validation to workgroup size (#20110)

Browse source 
### Description
add validation to workgroup size in `shaderHelper.mainStart()`.
This commit is contained in:
Yulong Wang 2024-04-02 19:29:20 -07:00 committed by GitHub
parent be831e1ba3
commit fa1917b81b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 18 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
js/web/lib/wasm/jsep/webgpu/ops/common.ts
View file

@ -783,7 +783,7 @@ export interface ShaderHelper {
}
class ShaderHelperImpl implements ShaderHelper {
constructor(private normalizedDispatchGroup: [number, number, number]) {}
constructor(private normalizedDispatchGroup: [number, number, number], private limits: GPUSupportedLimits) {}
guardAgainstOutOfBoundsWorkgroupSizes(size: number|string): string {
// Guard against out-of-bounds work group sizes
@ -796,6 +796,20 @@ class ShaderHelperImpl implements ShaderHelper {
const workgroupSizeY = typeof workgroupSize === 'number' ? 1 : workgroupSize[1];
const workgroupSizeZ = typeof workgroupSize === 'number' ? 1 : workgroupSize[2];
if (workgroupSizeX > this.limits.maxComputeWorkgroupSizeX ||
workgroupSizeY > this.limits.maxComputeWorkgroupSizeY ||
workgroupSizeZ > this.limits.maxComputeWorkgroupSizeZ) {
throw new Error(`workgroup size [${workgroupSizeX}, ${workgroupSizeY}, ${
workgroupSizeZ}] exceeds the maximum workgroup size [${this.limits.maxComputeWorkgroupSizeX}, ${
this.limits.maxComputeWorkgroupSizeY}, ${this.limits.maxComputeWorkgroupSizeZ}].`);
}
if (workgroupSizeX * workgroupSizeY * workgroupSizeZ > this.limits.maxComputeInvocationsPerWorkgroup) {
throw new Error(`workgroup size [${workgroupSizeX}, ${workgroupSizeY}, ${
workgroupSizeZ}] exceeds the maximum workgroup invocations ${
this.limits.maxComputeInvocationsPerWorkgroup}.`);
}
const is1DimensionDispatch = this.normalizedDispatchGroup[1] === 1 && this.normalizedDispatchGroup[2] === 1;
const paramList = is1DimensionDispatch ? `@builtin(global_invocation_id) global_id : vec3<u32>,
@builtin(workgroup_id) workgroup_id : vec3<u32>,
@ -920,7 +934,8 @@ class ShaderHelperImpl implements ShaderHelper {
}
}
export const createShaderHelper = (dispatchGroup: [number, number, number]) => new ShaderHelperImpl(dispatchGroup);
export const createShaderHelper = (dispatchGroup: [number, number, number], limits: GPUSupportedLimits) =>
new ShaderHelperImpl(dispatchGroup, limits);
/**
* This function comes from https://github.com/tensorflow/tfjs/blob/master/tfjs-core/src/ops/broadcast_util.ts#L18-L40

2
js/web/lib/wasm/jsep/webgpu/program-manager.ts
View file

@ -87,7 +87,7 @@ export class ProgramManager {
if (device.features.has('shader-f16')) {
extensions.push('enable f16;');
}
const shaderHelper = createShaderHelper(normalizedDispatchGroupSize);
const shaderHelper = createShaderHelper(normalizedDispatchGroupSize, this.backend.device.limits);
const userCode = programInfo.getShaderSource(shaderHelper);
const code = `${extensions.join('\n')}\n${shaderHelper.additionalImplementations}\n${userCode}`;
const shaderModule = device.createShaderModule({code, label: programInfo.name});