From eb92681bfbb95e93e34f301605529d605f7f9e63 Mon Sep 17 00:00:00 2001
From: Xu Xing <xing.xu@intel.com>
Date: Wed, 10 Jan 2024 01:33:57 +0800
Subject: [PATCH] [js/webgpu] Support range uniforms (#19055)

---
 js/web/lib/wasm/jsep/webgpu/ops/range.ts | 41 ++++++++++++++++--------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/js/web/lib/wasm/jsep/webgpu/ops/range.ts b/js/web/lib/wasm/jsep/webgpu/ops/range.ts
index 9cf66111bf..ed04b0f94b 100644
--- a/js/web/lib/wasm/jsep/webgpu/ops/range.ts
+++ b/js/web/lib/wasm/jsep/webgpu/ops/range.ts
@@ -3,10 +3,10 @@
 
 import {env} from 'onnxruntime-common';
 
-import {DataType} from '../../../wasm-common';
-import {ComputeContext, ProgramInfo} from '../types';
+import {DataType, tensorDataTypeEnumToString} from '../../../wasm-common';
+import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
 
-import {outputVariable, ShaderHelper} from './common';
+import {createTensorShapeVariables, outputVariable, ShaderHelper, UniformDataElementType, UniformsArrayType} from './common';
 
 const validateInputsContent = (start: number, limit: number, delta: number): void => {
   const sameStartLimit = start === limit;
@@ -22,23 +22,36 @@ const createRangeProgramInfo = (start: number, limit: number, delta: number, dat
   const numElements = Math.abs(Math.ceil((limit - start) / delta));
   const outputShape: number[] = [numElements];
   const outputSize = numElements;
+  const tensorDataType = tensorDataTypeEnumToString(dataType) as ProgramUniform['type'];
+  const programUniforms: ProgramUniform[] = [
+    {type: 'uint32', data: outputSize}, {type: tensorDataType, data: start}, {type: tensorDataType, data: delta},
+    ...createTensorShapeVariables(outputShape)
+  ];
 
-  const output = outputVariable('output', dataType, outputShape);
-  const wgslType = output.type.storage;
-
-  const getShaderSource = (shaderHelper: ShaderHelper) => `
-        ${shaderHelper.declareVariables(output)}
+  const getShaderSource = (shaderHelper: ShaderHelper) => {
+    const output = outputVariable('output', dataType, outputShape.length);
+    const wgslType = output.type.value;
+    const uniforms: UniformsArrayType = [
+      {name: 'outputSize', type: 'u32'}, {name: 'start', type: wgslType as UniformDataElementType},
+      {name: 'delta', type: wgslType as UniformDataElementType}
+    ];
+    return `
+        ${shaderHelper.registerUniforms(uniforms).declareVariables(output)}
         ${shaderHelper.mainStart()}
-        ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)}
-        output[global_idx] = ${wgslType}(${start}) + ${wgslType}(global_idx) * ${wgslType}(${delta});
+        ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')}
+        output[global_idx] = uniforms.start + ${wgslType}(global_idx) * uniforms.delta;
       }`;
+  };
+
   return {
     name: 'Range',
-    shaderCache: {hint: [start, limit, delta].map(x => x.toString()).join('_')},
+    shaderCache: {hint: `${dataType}`},
     getShaderSource,
-    getRunData: () => (
-        {outputs: [{dims: outputShape, dataType}],
-         dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}})
+    getRunData: () => ({
+      outputs: [{dims: outputShape, dataType}],
+      dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)},
+      programUniforms
+    })
   };
 };