From 49d7b3280119fcd1b828942e61413672ce2a882d Mon Sep 17 00:00:00 2001
From: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
Date: Thu, 12 Dec 2024 11:33:53 -0800
Subject: [PATCH] [js/webgpu] fix Conv2DMatMul shader's out-of-bound read
 (#23085)

### 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
---
 .../webgpu/ops/3rd-party/conv2d_mm_webgpu.ts     | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 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
index 3ef5c943d5624..9e21a552b8466 100644
--- 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,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);