From fa1917b81b8df9ea5befc0b256a715f717364c86 Mon Sep 17 00:00:00 2001
From: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
Date: Tue, 2 Apr 2024 19:29:20 -0700
Subject: [PATCH] [js/webgpu] add validation to workgroup size (#20110)

### Description
add validation to workgroup size in `shaderHelper.mainStart()`.
---
 js/web/lib/wasm/jsep/webgpu/ops/common.ts     | 19 +++++++++++++++++--
 .../lib/wasm/jsep/webgpu/program-manager.ts   |  2 +-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/js/web/lib/wasm/jsep/webgpu/ops/common.ts b/js/web/lib/wasm/jsep/webgpu/ops/common.ts
index f18ddb7206..17ac814c44 100644
--- a/js/web/lib/wasm/jsep/webgpu/ops/common.ts
+++ b/js/web/lib/wasm/jsep/webgpu/ops/common.ts
@@ -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
diff --git a/js/web/lib/wasm/jsep/webgpu/program-manager.ts b/js/web/lib/wasm/jsep/webgpu/program-manager.ts
index 0064e53a95..ccbcbe4850 100644
--- a/js/web/lib/wasm/jsep/webgpu/program-manager.ts
+++ b/js/web/lib/wasm/jsep/webgpu/program-manager.ts
@@ -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});