Skip to content

Commit

Permalink
[js/webgpu] fix Conv2DMatMul shader's out-of-bound read (#23085)
Browse files Browse the repository at this point in the history
### Description
<!-- Describe your changes. -->

Fix a bug caused by potential out-of-bound reads of `W` in the
Conv2DMatMul shader.

### Motivation and Context

Fixes #22983
  • Loading branch information
fs-eire authored and guschmue committed Dec 20, 2024
1 parent b9cf587 commit 49d7b32
Showing 1 changed file with 15 additions and 1 deletion.
16 changes: 15 additions & 1 deletion js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv2d_mm_webgpu.ts
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,21 @@ const conv2dCommonSnippet = (
}
return ${typeSnippet(innerElementSizeX, dataType)}(0.0);`;

const sampleW = `${getWSnippet(innerElementSizeW)}`;
const sampleW = isChannelsLast
? fitInner && fitBOuter
? getWSnippet(innerElementSizeW)
: `
let col = colIn * ${innerElementSizeW};
if (row < uniforms.dim_inner && col < uniforms.dim_b_outer) {
${getWSnippet(innerElementSizeW)}
}
return ${typeSnippet(innerElementSizeW, dataType)}(0.0);`
: `
let col = colIn * ${innerElementSizeW};
if (row < uniforms.dim_inner && col < uniforms.dim_a_outer) {
${getWSnippet(innerElementSizeW)}
}
return ${typeSnippet(innerElementSizeW, dataType)}(0.0);`;

const resType = typeSnippet(innerElementSize, dataType);
const aType = isChannelsLast ? typeSnippet(innerElementSizeX, dataType) : typeSnippet(innerElementSizeW, dataType);
Expand Down

0 comments on commit 49d7b32

Please sign in to comment.