From 645802f6a7fe7c74696ec62a350c0e4c902b1183 Mon Sep 17 00:00:00 2001 From: Linchenn Date: Wed, 23 Aug 2023 15:46:47 -0700 Subject: [PATCH 1/7] Add uniforms --- tfjs-backend-webgpu/src/backend_webgpu.ts | 6 +++--- tfjs-backend-webgpu/src/webgpu_program.ts | 16 ++++++++++++---- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts index ba3ffecd2a6..1b8da1e7ee7 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu.ts @@ -60,7 +60,7 @@ export interface WebGPUTimingInfo extends TimingInfo { downloadWaitMs: number; } -type ProgramUniform = Array<{type: string; data: number[]}>; +export type ProgramUniform = Array<{type: string; data: number[]}>; // Empirically determined constant used to determine size threshold for handing // off execution to the CPU. @@ -882,8 +882,8 @@ export class WebGPUBackend extends KernelBackend { }; }); - program.shaderKey = - webgpu_program.makeShaderKey(program, inputsData, output); + program.shaderKey = webgpu_program.makeShaderKey( + program, inputsData, output, programDefinedUniform); const parallelCompilation = env().getBool('WEBGPU_ENGINE_COMPILE_ONLY'); if (!(program.shaderKey in this.pipelineCache)) { diff --git a/tfjs-backend-webgpu/src/webgpu_program.ts b/tfjs-backend-webgpu/src/webgpu_program.ts index ddaeddc6b9a..f2fe599c344 100644 --- a/tfjs-backend-webgpu/src/webgpu_program.ts +++ b/tfjs-backend-webgpu/src/webgpu_program.ts @@ -17,6 +17,7 @@ import {backend_util, DataType, DataTypeMap, env, Rank, TensorInfo, util} from '@tensorflow/tfjs-core'; +import {ProgramUniform} from './backend_webgpu'; import {symbolicallyComputeStrides} from './shader_util'; export enum PixelsOpType { @@ -348,8 +349,8 @@ function makeShader( } export function makeShaderKey( - program: WebGPUProgram, inputsData: InputInfo[], - output: TensorInfo): string { + program: WebGPUProgram, inputsData: InputInfo[], output: TensorInfo, + programDefinedUniform?: ProgramUniform): string { let key = program.shaderKey; if (program.pixelsOpType != null) { return key; @@ -361,6 +362,13 @@ export function makeShaderKey( shapes.push(element.shape); types.push(element.dtype); }); + const uniformLengths: number[] = []; + if (programDefinedUniform != null) { + programDefinedUniform.forEach(element => { + uniformLengths.push(element.data.length); + types.push(element.type); + }); + } shapes.push(output.shape); types.push(output.dtype); @@ -373,8 +381,8 @@ export function makeShaderKey( const flatDispatchString = isFlatDispatch(program) ? 'flatDispatch' : ''; key += '_' + (program.workgroupSize ? program.workgroupSize.join(',') : '') + - shapes.map(shape => shape.length).join(',') + types.join(',') + - program.variableNames.join(',') + broadcastDimsKey + + shapes.map(shape => shape.length).join(',') + uniformLengths.join(',') + + types.join(',') + program.variableNames.join(',') + broadcastDimsKey + inputShapesEqualsOutShape + flatDispatchString; return key; From 58ea96fc1b3b1bbd53b792499d1946d5aa367a28 Mon Sep 17 00:00:00 2001 From: Linchenn Date: Wed, 23 Aug 2023 16:04:26 -0700 Subject: [PATCH 2/7] Update webgpu_program.ts --- tfjs-backend-webgpu/src/webgpu_program.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgpu/src/webgpu_program.ts b/tfjs-backend-webgpu/src/webgpu_program.ts index f2fe599c344..ee88b6335e9 100644 --- a/tfjs-backend-webgpu/src/webgpu_program.ts +++ b/tfjs-backend-webgpu/src/webgpu_program.ts @@ -357,7 +357,7 @@ export function makeShaderKey( } const shapes: number[][] = []; - const types: Array = []; + const types: Array = []; inputsData.forEach(element => { shapes.push(element.shape); types.push(element.dtype); From 32386acc15f15293987b9ce75dadf71775f6a01c Mon Sep 17 00:00:00 2001 From: Linchenn Date: Wed, 23 Aug 2023 16:54:24 -0700 Subject: [PATCH 3/7] Update webgpu_program.ts --- tfjs-backend-webgpu/src/webgpu_program.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgpu/src/webgpu_program.ts b/tfjs-backend-webgpu/src/webgpu_program.ts index ee88b6335e9..66e6f5044c0 100644 --- a/tfjs-backend-webgpu/src/webgpu_program.ts +++ b/tfjs-backend-webgpu/src/webgpu_program.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -import {backend_util, DataType, DataTypeMap, env, Rank, TensorInfo, util} from '@tensorflow/tfjs-core'; +import {backend_util, DataType, env, Rank, TensorInfo, util} from '@tensorflow/tfjs-core'; import {ProgramUniform} from './backend_webgpu'; import {symbolicallyComputeStrides} from './shader_util'; From bad5fb99ac64b8afbe2fc272c84a99812ac41369 Mon Sep 17 00:00:00 2001 From: Linchenn Date: Tue, 19 Sep 2023 15:50:06 -0700 Subject: [PATCH 4/7] Revert "Add uniforms" This reverts commit 645802f6a7fe7c74696ec62a350c0e4c902b1183. --- tfjs-backend-webgpu/src/backend_webgpu.ts | 6 +++--- tfjs-backend-webgpu/src/webgpu_program.ts | 16 ++++------------ 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts index 1b8da1e7ee7..ba3ffecd2a6 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu.ts @@ -60,7 +60,7 @@ export interface WebGPUTimingInfo extends TimingInfo { downloadWaitMs: number; } -export type ProgramUniform = Array<{type: string; data: number[]}>; +type ProgramUniform = Array<{type: string; data: number[]}>; // Empirically determined constant used to determine size threshold for handing // off execution to the CPU. @@ -882,8 +882,8 @@ export class WebGPUBackend extends KernelBackend { }; }); - program.shaderKey = webgpu_program.makeShaderKey( - program, inputsData, output, programDefinedUniform); + program.shaderKey = + webgpu_program.makeShaderKey(program, inputsData, output); const parallelCompilation = env().getBool('WEBGPU_ENGINE_COMPILE_ONLY'); if (!(program.shaderKey in this.pipelineCache)) { diff --git a/tfjs-backend-webgpu/src/webgpu_program.ts b/tfjs-backend-webgpu/src/webgpu_program.ts index 66e6f5044c0..d7b15bbdde3 100644 --- a/tfjs-backend-webgpu/src/webgpu_program.ts +++ b/tfjs-backend-webgpu/src/webgpu_program.ts @@ -17,7 +17,6 @@ import {backend_util, DataType, env, Rank, TensorInfo, util} from '@tensorflow/tfjs-core'; -import {ProgramUniform} from './backend_webgpu'; import {symbolicallyComputeStrides} from './shader_util'; export enum PixelsOpType { @@ -349,8 +348,8 @@ function makeShader( } export function makeShaderKey( - program: WebGPUProgram, inputsData: InputInfo[], output: TensorInfo, - programDefinedUniform?: ProgramUniform): string { + program: WebGPUProgram, inputsData: InputInfo[], + output: TensorInfo): string { let key = program.shaderKey; if (program.pixelsOpType != null) { return key; @@ -362,13 +361,6 @@ export function makeShaderKey( shapes.push(element.shape); types.push(element.dtype); }); - const uniformLengths: number[] = []; - if (programDefinedUniform != null) { - programDefinedUniform.forEach(element => { - uniformLengths.push(element.data.length); - types.push(element.type); - }); - } shapes.push(output.shape); types.push(output.dtype); @@ -381,8 +373,8 @@ export function makeShaderKey( const flatDispatchString = isFlatDispatch(program) ? 'flatDispatch' : ''; key += '_' + (program.workgroupSize ? program.workgroupSize.join(',') : '') + - shapes.map(shape => shape.length).join(',') + uniformLengths.join(',') + - types.join(',') + program.variableNames.join(',') + broadcastDimsKey + + shapes.map(shape => shape.length).join(',') + types.join(',') + + program.variableNames.join(',') + broadcastDimsKey + inputShapesEqualsOutShape + flatDispatchString; return key; From 33d9ad41a1bb0b0df78042824d438f19a5fd886d Mon Sep 17 00:00:00 2001 From: Linchenn Date: Tue, 19 Sep 2023 15:50:13 -0700 Subject: [PATCH 5/7] Revert "Update webgpu_program.ts" This reverts commit 58ea96fc1b3b1bbd53b792499d1946d5aa367a28. --- tfjs-backend-webgpu/src/webgpu_program.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgpu/src/webgpu_program.ts b/tfjs-backend-webgpu/src/webgpu_program.ts index d7b15bbdde3..b0beb1e3197 100644 --- a/tfjs-backend-webgpu/src/webgpu_program.ts +++ b/tfjs-backend-webgpu/src/webgpu_program.ts @@ -356,7 +356,7 @@ export function makeShaderKey( } const shapes: number[][] = []; - const types: Array = []; + const types: Array = []; inputsData.forEach(element => { shapes.push(element.shape); types.push(element.dtype); From 3341e43134aeeb2da922c015103ca7304905f032 Mon Sep 17 00:00:00 2001 From: Linchenn Date: Tue, 19 Sep 2023 15:50:19 -0700 Subject: [PATCH 6/7] Revert "Update webgpu_program.ts" This reverts commit 32386acc15f15293987b9ce75dadf71775f6a01c. --- tfjs-backend-webgpu/src/webgpu_program.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgpu/src/webgpu_program.ts b/tfjs-backend-webgpu/src/webgpu_program.ts index b0beb1e3197..ddaeddc6b9a 100644 --- a/tfjs-backend-webgpu/src/webgpu_program.ts +++ b/tfjs-backend-webgpu/src/webgpu_program.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -import {backend_util, DataType, env, Rank, TensorInfo, util} from '@tensorflow/tfjs-core'; +import {backend_util, DataType, DataTypeMap, env, Rank, TensorInfo, util} from '@tensorflow/tfjs-core'; import {symbolicallyComputeStrides} from './shader_util'; From 6ac1472eb90faf4c109d90db1887bdd35cec3717 Mon Sep 17 00:00:00 2001 From: Linchenn Date: Tue, 19 Sep 2023 15:51:49 -0700 Subject: [PATCH 7/7] Add key to scatter webgpu program --- tfjs-backend-webgpu/src/scatter_webgpu.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tfjs-backend-webgpu/src/scatter_webgpu.ts b/tfjs-backend-webgpu/src/scatter_webgpu.ts index e866424a33a..22402030387 100644 --- a/tfjs-backend-webgpu/src/scatter_webgpu.ts +++ b/tfjs-backend-webgpu/src/scatter_webgpu.ts @@ -48,8 +48,9 @@ export class ScatterProgram implements WebGPUProgram { this.dispatch = computeDispatch(this.dispatchLayout, flattenXShape, this.workgroupSize); this.sliceDimGreaterThanOne = sliceDim > 1; - this.shaderKey = `scatter_${indicesRank}_${updatesRank}_${ - this.sliceDimGreaterThanOne}_${outputDtype}_${sumDupeIndices}`; + this.shaderKey = + `scatter_${indicesRank}_${updatesRank}_${this.sliceDimGreaterThanOne}_${ + outputDtype}_${sumDupeIndices}_${strides.length}`; const stridesType = getCoordsDataType(strides.length); this.uniforms = `sliceDim : i32, strides: ${stridesType}, updatesSize: i32,`;