[js/webgpu] fix Conv2DMatMul shader's out-of-bound read

microsoft · Dec 11, 2024 · 943e2b5 · 943e2b5
1 parent 2f2c73b
commit 943e2b5
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv2d_mm_webgpu.ts b/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv2d_mm_webgpu.ts
@@ -143,7 +143,15 @@ const conv2dCommonSnippet = (
     }
     return ${typeSnippet(innerElementSizeX, dataType)}(0.0);`;
 
-  const sampleW = `${getWSnippet(innerElementSizeW)}`;
+  const sampleW =
+    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);`;
 
   const resType = typeSnippet(innerElementSize, dataType);
   const aType = isChannelsLast ? typeSnippet(innerElementSizeX, dataType) : typeSnippet(innerElementSizeW, dataType);