From f4a22bfd5212edc3dfe8f4c8a5fa1395d48638a7 Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Wed, 21 Feb 2024 13:27:18 -0800 Subject: [PATCH 1/7] Revert "[JS/WebGPU] Added uniforms to Tile and Where Ops (#18768)" This reverts commit d673e39ad89a709d5896510bcd496927567b4b79. --- js/web/lib/wasm/jsep/webgpu/ops/tile.ts | 25 +++++------ js/web/lib/wasm/jsep/webgpu/ops/where.ts | 56 ++++++++++++------------ 2 files changed, 39 insertions(+), 42 deletions(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/tile.ts b/js/web/lib/wasm/jsep/webgpu/ops/tile.ts index f9728575fe072..e294541a775ca 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/tile.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/tile.ts @@ -6,7 +6,7 @@ import {TensorView} from '../../tensor-view'; import {ShapeUtil} from '../../util'; import {ComputeContext, ProgramInfo} from '../types'; -import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common'; +import {inputVariable, outputVariable, ShaderHelper} from './common'; const getRepeats = (repeatsTensorView: TensorView): readonly number[] => Array.from(repeatsTensorView.getBigInt64Array(), Number); @@ -54,33 +54,30 @@ export const createTileProgramInfo = (inputs: readonly TensorView[]): ProgramInf const outputSize = ShapeUtil.size(outputShape); const dataType = inputs[0].dataType; - const input = inputVariable('input', dataType, inputShape.length); - const output = outputVariable('output', dataType, outputShape.length); + const input = inputVariable('input', dataType, inputShape); + const output = outputVariable('output', dataType, outputShape); const getShaderSource = (shaderHelper: ShaderHelper) => ` const inputShape = ${input.indices(...inputShape)}; - ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)} + ${shaderHelper.declareVariables(input, output)} ${shaderHelper.mainStart()} - ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.output_size')} - let output_indices = ${output.offsetToIndices('global_idx')}; - var input_indices: ${input.type.indices}; + ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)} + let outputIndices = ${output.offsetToIndices('global_idx')}; + var inputIndices: ${input.type.indices}; for (var i = 0; i < ${inputShape.length}; i++) { - let input_dim_i = ${input.indicesGet('uniforms.input_shape', 'i')}; - let input_dim_value = ${output.indicesGet('output_indices', 'i')} % input_dim_i; + let inputDimValue = ${output.indicesGet('outputIndices', 'i')} % ${input.indicesGet('inputShape', 'i')}; - ${input.indicesSet('input_indices', 'i', 'input_dim_value')} + ${input.indicesSet('inputIndices', 'i', 'inputDimValue')} } - ${output.setByOffset('global_idx', input.getByIndices('input_indices'))} + ${output.setByOffset('global_idx', input.getByIndices('inputIndices'))} }`; return { name: 'Tile', - shaderCache: {hint: `${repeats}`, inputDependencies: ['rank']}, + shaderCache: {hint: `${repeats}`}, getRunData: () => ({ outputs: [{dims: outputShape, dataType: inputs[0].dataType}], dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}, - programUniforms: - [{type: DataType.uint32, data: outputSize}, ...createTensorShapeVariables(inputs[0].dims, outputShape)], }), getShaderSource, }; diff --git a/js/web/lib/wasm/jsep/webgpu/ops/where.ts b/js/web/lib/wasm/jsep/webgpu/ops/where.ts index cfee07a9239d7..6e7a055187754 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/where.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/where.ts @@ -6,15 +6,18 @@ import {TensorView} from '../../tensor-view'; import {BroadcastUtil, ShapeUtil} from '../../util'; import {ComputeContext, ProgramInfo} from '../types'; -import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common'; +import {inputVariable, outputVariable, ShaderHelper} from './common'; const createWhereOpProgramShader = (shaderHelper: ShaderHelper, inputs: readonly TensorView[], dimsOutput: readonly number[], isBroadcast: boolean, typeOutput: number) => { - const output = outputVariable('output_data', typeOutput, dimsOutput.length, 4); - const a = inputVariable('a_data', inputs[1].dataType, inputs[1].dims.length, 4); - const b = inputVariable('b_data', inputs[2].dataType, inputs[2].dims.length, 4); - const c = inputVariable('c_data', inputs[0].dataType, inputs[0].dims.length, 4); + const outputSize = ShapeUtil.size(dimsOutput); + const vecSize = Math.ceil(outputSize / 4); + + const output = outputVariable('outputData', typeOutput, dimsOutput, 4); + const a = inputVariable('aData', inputs[1].dataType, inputs[1].dims, 4); + const b = inputVariable('bData', inputs[2].dataType, inputs[2].dims, 4); + const c = inputVariable('cData', inputs[0].dataType, inputs[0].dims, 4); let assignment: string; const expression = (a: string, b: string, c: string) => `select(${b}, ${a}, ${c})`; @@ -24,20 +27,20 @@ const createWhereOpProgramShader = expression(a.getByOffset('global_idx'), b.getByOffset('global_idx'), c.getByOffset('global_idx'))); } else { const singleAssignment = (resStr: string, x: number, typeCast = '') => { - const expressionA = `a_data[index_a${x}][component_a${x}]`; - const expressionB = `b_data[index_b${x}][component_b${x}]`; + const expressionA = `aData[indexA${x}][componentA${x}]`; + const expressionB = `bData[indexB${x}][componentB${x}]`; // eslint-disable-next-line no-bitwise - const expressionC = `bool(c_data[index_c${x}] & ${0xff000000 >>> ((3 - x) * 8)}u)`; + const expressionC = `bool(cData[indexC${x}] & ${0xff000000 >>> ((3 - x) * 8)}u)`; return ` - let output_indices${x} = ${output.offsetToIndices(`global_idx * 4u + ${x}u`)}; - let offset_a${x} = ${a.broadcastedIndicesToOffset(`output_indices${x}`, output)}; - let offset_b${x} = ${b.broadcastedIndicesToOffset(`output_indices${x}`, output)}; - let offset_c${x} = ${c.broadcastedIndicesToOffset(`output_indices${x}`, output)}; - let index_a${x} = offset_a${x} / 4u; - let index_b${x} = offset_b${x} / 4u; - let index_c${x} = offset_c${x} / 4u; - let component_a${x} = offset_a${x} % 4u; - let component_b${x} = offset_b${x} % 4u; + let outputIndices${x} = ${output.offsetToIndices(`global_idx * 4u + ${x}u`)}; + let offsetA${x} = ${a.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; + let offsetB${x} = ${b.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; + let offsetC${x} = ${c.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; + let indexA${x} = offsetA${x} / 4u; + let indexB${x} = offsetB${x} / 4u; + let indexC${x} = offsetC${x} / 4u; + let componentA${x} = offsetA${x} % 4u; + let componentB${x} = offsetB${x} % 4u; ${resStr}[${x}] = ${typeCast}(${expression(expressionA, expressionB, expressionC)}); `; }; @@ -48,21 +51,21 @@ const createWhereOpProgramShader = ${singleAssignment('data', 1, 'u32')} ${singleAssignment('data', 2, 'u32')} ${singleAssignment('data', 3, 'u32')} - output_data[global_idx] = dot(vec4(0x1, 0x100, 0x10000, 0x1000000), vec4(data));`; + outputData[global_idx] = dot(vec4(0x1, 0x100, 0x10000, 0x1000000), vec4(data));`; } else { assignment = ` - ${singleAssignment('output_data[global_idx]', 0)} - ${singleAssignment('output_data[global_idx]', 1)} - ${singleAssignment('output_data[global_idx]', 2)} - ${singleAssignment('output_data[global_idx]', 3)} + ${singleAssignment('outputData[global_idx]', 0)} + ${singleAssignment('outputData[global_idx]', 1)} + ${singleAssignment('outputData[global_idx]', 2)} + ${singleAssignment('outputData[global_idx]', 3)} `; } } return ` - ${shaderHelper.registerUniform('vec_size', 'u32').declareVariables(c, a, b, output)} + ${shaderHelper.declareVariables(c, a, b, output)} ${shaderHelper.mainStart()} - ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.vec_size')} + ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(vecSize)} ${assignment} }`; }; @@ -91,14 +94,11 @@ const createWhereOpProgramInfo = (inputs: readonly TensorView[]): ProgramInfo => return { name: 'Where', - shaderCache: {inputDependencies: ['rank', 'rank', 'rank']}, getShaderSource: (shaderHelper) => createWhereOpProgramShader(shaderHelper, inputs, outputShape, isBroadcast, outputDataType), getRunData: () => ({ outputs: [{dims: outputShape, dataType: outputDataType}], - dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* vec size */)}, - programUniforms: - [{type: DataType.uint32, data: vecSize}, ...createTensorShapeVariables(dimsC, dimsA, dimsB, outputShape)], + dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* vec size */)} }), }; }; From c581bd86e7c1f6de85325d642e30230e2142eafa Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Thu, 22 Feb 2024 08:38:14 -0800 Subject: [PATCH 2/7] Added inputShape to shaderChche.hint in Split implementation. --- js/web/lib/wasm/jsep/webgpu/ops/split.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/split.ts b/js/web/lib/wasm/jsep/webgpu/ops/split.ts index a09ac78b17006..701706ddcdaca 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/split.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/split.ts @@ -107,7 +107,7 @@ const createSplitProgramInfo = (inputs: readonly TensorView[], attributes: Split }`; return { name: 'Split', - shaderCache: {hint: attributes.cacheKey, inputDependencies: ['rank']}, + shaderCache: {hint: `${attributes.cacheKey};${inputShape}`, inputDependencies: ['rank']}, getShaderSource, getRunData: () => ({ outputs: outputsTensorInfo, From f9e7e4fa77683c0bee3949c7df74e68d4a182436 Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Thu, 22 Feb 2024 12:17:52 -0800 Subject: [PATCH 3/7] Fixed an issue with packing 4 boolean input into one uint32 --- js/web/lib/wasm/jsep/webgpu/ops/where.ts | 5 ++-- js/web/test/data/ops/where.jsonc | 34 ++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/where.ts b/js/web/lib/wasm/jsep/webgpu/ops/where.ts index 6e7a055187754..e3970233ac8cb 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/where.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/where.ts @@ -30,7 +30,7 @@ const createWhereOpProgramShader = const expressionA = `aData[indexA${x}][componentA${x}]`; const expressionB = `bData[indexB${x}][componentB${x}]`; // eslint-disable-next-line no-bitwise - const expressionC = `bool(cData[indexC${x}] & ${0xff000000 >>> ((3 - x) * 8)}u)`; + const expressionC = `bool(cData[indexC${x}] & (0xffu << (componentC${x} * 8)))`; return ` let outputIndices${x} = ${output.offsetToIndices(`global_idx * 4u + ${x}u`)}; let offsetA${x} = ${a.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; @@ -41,6 +41,7 @@ const createWhereOpProgramShader = let indexC${x} = offsetC${x} / 4u; let componentA${x} = offsetA${x} % 4u; let componentB${x} = offsetB${x} % 4u; + let componentC${x} = offsetC${x} % 4u; ${resStr}[${x}] = ${typeCast}(${expression(expressionA, expressionB, expressionC)}); `; }; @@ -90,8 +91,6 @@ const createWhereOpProgramInfo = (inputs: readonly TensorView[]): ProgramInfo => outputSize = ShapeUtil.size(outputShape); } - const vecSize = Math.ceil(outputSize / 4); - return { name: 'Where', getShaderSource: (shaderHelper) => diff --git a/js/web/test/data/ops/where.jsonc b/js/web/test/data/ops/where.jsonc index 047fd6fd7511b..990120dd3708e 100644 --- a/js/web/test/data/ops/where.jsonc +++ b/js/web/test/data/ops/where.jsonc @@ -168,5 +168,39 @@ ] } ] + }, + { + "name": "Where with no attributes", + "operator": "Where", + "attributes": [], + "cases": [ + { + "name": "T[1 1 2 1] T[1 4] T[1 1 2 4] float32 broadcast 1", + "inputs": [ + { + "data": [true, false], + "dims": [1, 1, 2, 1], + "type": "bool" + }, + { + "data": [1, 2, 3, 4], + "dims": [1, 4], + "type": "float32" + }, + { + "data": [5, 6, 7, 8, 9, 10, 11, 12], + "dims": [1, 1, 2, 4], + "type": "float32" + } + ], + "outputs": [ + { + "data": [1, 2, 3, 4, 9, 10, 11, 12], + "dims": [1, 1, 2, 4], + "type": "float32" + } + ] + } + ] } ] From f0b745171473f2e83b67661b082a7a48ad7f4bc9 Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Thu, 22 Feb 2024 12:18:33 -0800 Subject: [PATCH 4/7] Fixed an error in referencing input tensors. --- js/web/test/test-runner.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/js/web/test/test-runner.ts b/js/web/test/test-runner.ts index ecc7d4b4a09a5..a4adf5c4ce144 100644 --- a/js/web/test/test-runner.ts +++ b/js/web/test/test-runner.ts @@ -627,8 +627,8 @@ export async function runModelTestSet( try { const feeds: Record = {}; const outputsMetaInfo: Record = {}; - testCase.inputs!.forEach((tensor, i) => feeds[context.session.inputNames[i]] = tensor); - testCase.outputs!.forEach((tensor, i) => outputsMetaInfo[context.session.outputNames[i]] = tensor); + testCase.inputs!.forEach((tensor) => feeds[tensor.name] = tensor); + testCase.outputs!.forEach((tensor) => outputsMetaInfo[tensor.name] = tensor); const [start, end, outputs] = await sessionRun({session: context.session, feeds, outputsMetaInfo, ioBinding: context.ioBinding}); if (context.perfData.count === 0) { From 4d06b604faf7617db6835a1f7960e7d2c30afc98 Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Thu, 22 Feb 2024 16:42:01 -0800 Subject: [PATCH 5/7] Revert unintended changes to tile.js. --- js/web/lib/wasm/jsep/webgpu/ops/tile.ts | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/tile.ts b/js/web/lib/wasm/jsep/webgpu/ops/tile.ts index e294541a775ca..f9728575fe072 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/tile.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/tile.ts @@ -6,7 +6,7 @@ import {TensorView} from '../../tensor-view'; import {ShapeUtil} from '../../util'; import {ComputeContext, ProgramInfo} from '../types'; -import {inputVariable, outputVariable, ShaderHelper} from './common'; +import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common'; const getRepeats = (repeatsTensorView: TensorView): readonly number[] => Array.from(repeatsTensorView.getBigInt64Array(), Number); @@ -54,30 +54,33 @@ export const createTileProgramInfo = (inputs: readonly TensorView[]): ProgramInf const outputSize = ShapeUtil.size(outputShape); const dataType = inputs[0].dataType; - const input = inputVariable('input', dataType, inputShape); - const output = outputVariable('output', dataType, outputShape); + const input = inputVariable('input', dataType, inputShape.length); + const output = outputVariable('output', dataType, outputShape.length); const getShaderSource = (shaderHelper: ShaderHelper) => ` const inputShape = ${input.indices(...inputShape)}; - ${shaderHelper.declareVariables(input, output)} + ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)} ${shaderHelper.mainStart()} - ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)} - let outputIndices = ${output.offsetToIndices('global_idx')}; - var inputIndices: ${input.type.indices}; + ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.output_size')} + let output_indices = ${output.offsetToIndices('global_idx')}; + var input_indices: ${input.type.indices}; for (var i = 0; i < ${inputShape.length}; i++) { - let inputDimValue = ${output.indicesGet('outputIndices', 'i')} % ${input.indicesGet('inputShape', 'i')}; + let input_dim_i = ${input.indicesGet('uniforms.input_shape', 'i')}; + let input_dim_value = ${output.indicesGet('output_indices', 'i')} % input_dim_i; - ${input.indicesSet('inputIndices', 'i', 'inputDimValue')} + ${input.indicesSet('input_indices', 'i', 'input_dim_value')} } - ${output.setByOffset('global_idx', input.getByIndices('inputIndices'))} + ${output.setByOffset('global_idx', input.getByIndices('input_indices'))} }`; return { name: 'Tile', - shaderCache: {hint: `${repeats}`}, + shaderCache: {hint: `${repeats}`, inputDependencies: ['rank']}, getRunData: () => ({ outputs: [{dims: outputShape, dataType: inputs[0].dataType}], dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}, + programUniforms: + [{type: DataType.uint32, data: outputSize}, ...createTensorShapeVariables(inputs[0].dims, outputShape)], }), getShaderSource, }; From f0205173f8e7107bdcad0131841b98556da141db Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Thu, 22 Feb 2024 16:50:37 -0800 Subject: [PATCH 6/7] Revert unintended changes to where.js --- js/web/lib/wasm/jsep/webgpu/ops/where.ts | 60 ++++++++++++------------ 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/where.ts b/js/web/lib/wasm/jsep/webgpu/ops/where.ts index e3970233ac8cb..a6375847fc42f 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/where.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/where.ts @@ -6,18 +6,15 @@ import {TensorView} from '../../tensor-view'; import {BroadcastUtil, ShapeUtil} from '../../util'; import {ComputeContext, ProgramInfo} from '../types'; -import {inputVariable, outputVariable, ShaderHelper} from './common'; +import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common'; const createWhereOpProgramShader = (shaderHelper: ShaderHelper, inputs: readonly TensorView[], dimsOutput: readonly number[], isBroadcast: boolean, typeOutput: number) => { - const outputSize = ShapeUtil.size(dimsOutput); - const vecSize = Math.ceil(outputSize / 4); - - const output = outputVariable('outputData', typeOutput, dimsOutput, 4); - const a = inputVariable('aData', inputs[1].dataType, inputs[1].dims, 4); - const b = inputVariable('bData', inputs[2].dataType, inputs[2].dims, 4); - const c = inputVariable('cData', inputs[0].dataType, inputs[0].dims, 4); + const output = outputVariable('output_data', typeOutput, dimsOutput.length, 4); + const a = inputVariable('a_data', inputs[1].dataType, inputs[1].dims.length, 4); + const b = inputVariable('b_data', inputs[2].dataType, inputs[2].dims.length, 4); + const c = inputVariable('c_data', inputs[0].dataType, inputs[0].dims.length, 4); let assignment: string; const expression = (a: string, b: string, c: string) => `select(${b}, ${a}, ${c})`; @@ -27,21 +24,21 @@ const createWhereOpProgramShader = expression(a.getByOffset('global_idx'), b.getByOffset('global_idx'), c.getByOffset('global_idx'))); } else { const singleAssignment = (resStr: string, x: number, typeCast = '') => { - const expressionA = `aData[indexA${x}][componentA${x}]`; - const expressionB = `bData[indexB${x}][componentB${x}]`; + const expressionA = `a_data[index_a${x}][component_a${x}]`; + const expressionB = `b_data[index_b${x}][component_b${x}]`; // eslint-disable-next-line no-bitwise - const expressionC = `bool(cData[indexC${x}] & (0xffu << (componentC${x} * 8)))`; + const expressionC = `bool(c_data[index_c${x}] & (0xffu << (component_c${x} * 8)))`; return ` - let outputIndices${x} = ${output.offsetToIndices(`global_idx * 4u + ${x}u`)}; - let offsetA${x} = ${a.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; - let offsetB${x} = ${b.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; - let offsetC${x} = ${c.broadcastedIndicesToOffset(`outputIndices${x}`, output)}; - let indexA${x} = offsetA${x} / 4u; - let indexB${x} = offsetB${x} / 4u; - let indexC${x} = offsetC${x} / 4u; - let componentA${x} = offsetA${x} % 4u; - let componentB${x} = offsetB${x} % 4u; - let componentC${x} = offsetC${x} % 4u; + let output_indices${x} = ${output.offsetToIndices(`global_idx * 4u + ${x}u`)}; + let offset_a${x} = ${a.broadcastedIndicesToOffset(`output_indices${x}`, output)}; + let offset_b${x} = ${b.broadcastedIndicesToOffset(`output_indices${x}`, output)}; + let offset_c${x} = ${c.broadcastedIndicesToOffset(`output_indices${x}`, output)}; + let index_a${x} = offset_a${x} / 4u; + let index_b${x} = offset_b${x} / 4u; + let index_c${x} = offset_c${x} / 4u; + let component_a${x} = offset_a${x} % 4u; + let component_b${x} = offset_b${x} % 4u; + let component_c${x} = offset_c${x} % 4u; ${resStr}[${x}] = ${typeCast}(${expression(expressionA, expressionB, expressionC)}); `; }; @@ -52,21 +49,21 @@ const createWhereOpProgramShader = ${singleAssignment('data', 1, 'u32')} ${singleAssignment('data', 2, 'u32')} ${singleAssignment('data', 3, 'u32')} - outputData[global_idx] = dot(vec4(0x1, 0x100, 0x10000, 0x1000000), vec4(data));`; + output_data[global_idx] = dot(vec4(0x1, 0x100, 0x10000, 0x1000000), vec4(data));`; } else { assignment = ` - ${singleAssignment('outputData[global_idx]', 0)} - ${singleAssignment('outputData[global_idx]', 1)} - ${singleAssignment('outputData[global_idx]', 2)} - ${singleAssignment('outputData[global_idx]', 3)} + ${singleAssignment('output_data[global_idx]', 0)} + ${singleAssignment('output_data[global_idx]', 1)} + ${singleAssignment('output_data[global_idx]', 2)} + ${singleAssignment('output_data[global_idx]', 3)} `; } } return ` - ${shaderHelper.declareVariables(c, a, b, output)} + ${shaderHelper.registerUniform('vec_size', 'u32').declareVariables(c, a, b, output)} ${shaderHelper.mainStart()} - ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(vecSize)} + ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.vec_size')} ${assignment} }`; }; @@ -91,13 +88,18 @@ const createWhereOpProgramInfo = (inputs: readonly TensorView[]): ProgramInfo => outputSize = ShapeUtil.size(outputShape); } + const vecSize = Math.ceil(outputSize / 4); + return { name: 'Where', + shaderCache: {inputDependencies: ['rank', 'rank', 'rank']}, getShaderSource: (shaderHelper) => createWhereOpProgramShader(shaderHelper, inputs, outputShape, isBroadcast, outputDataType), getRunData: () => ({ outputs: [{dims: outputShape, dataType: outputDataType}], - dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* vec size */)} + dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* vec size */)}, + programUniforms: + [{type: DataType.uint32, data: vecSize}, ...createTensorShapeVariables(dimsC, dimsA, dimsB, outputShape)], }), }; }; From 47eb6eedf6369b21614b174d5f64f1e18debe629 Mon Sep 17 00:00:00 2001 From: Satya Jandhyala Date: Thu, 22 Feb 2024 19:31:19 -0800 Subject: [PATCH 7/7] Reverted changes to split.ts --- js/web/lib/wasm/jsep/webgpu/ops/split.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/split.ts b/js/web/lib/wasm/jsep/webgpu/ops/split.ts index 701706ddcdaca..a09ac78b17006 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/split.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/split.ts @@ -107,7 +107,7 @@ const createSplitProgramInfo = (inputs: readonly TensorView[], attributes: Split }`; return { name: 'Split', - shaderCache: {hint: `${attributes.cacheKey};${inputShape}`, inputDependencies: ['rank']}, + shaderCache: {hint: attributes.cacheKey, inputDependencies: ['rank']}, getShaderSource, getRunData: () => ({ outputs: outputsTensorInfo,