Skip to content

[webgpu] fix shaderKey issue to add types info in shaderKey #3049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
annxingyuan merged 4 commits into from
Apr 18, 2020
Merged

[webgpu] fix shaderKey issue to add types info in shaderKey #3049

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
61b3f85
[webgpu] fix shaderKey issue to add types info in shaderKey
haoyunfeix Apr 8, 2020
304923b
Merge branch 'master' into fix_shader_key
haoyunfeix Apr 16, 2020
3b98260
Merge branch 'master' into fix_shader_key
annxingyuan Apr 18, 2020
90858f5
Merge branch 'master' into fix_shader_key
annxingyuan Apr 18, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions tfjs-backend-webgpu/src/backend_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -457,8 +457,6 @@ export class WebGPUBackend extends KernelBackend {
const uniformData = new Int32Array(dimUniforms);
const uniforms = this.makeUniforms(uniformData);

const key =
webgpu_program.makeShaderKey(program, bufferShapes.map(d => d.length));
const inputsData = inputs.map((input: Tensor, i: number) => {
this.uploadToGPU(input.dataId);

Expand All @@ -471,6 +469,9 @@ export class WebGPUBackend extends KernelBackend {
};
});
this.uploadToGPU(output.dataId);
const bufferTypes = inputsData.map(d => d.dtype).concat(output.dtype);
const key =
webgpu_program.makeShaderKey(program, bufferShapes, bufferTypes);
const {bindGroupLayout, pipeline} = this.getAndSavePipeline(key, () => {
return webgpu_program.compileProgram(
this.glslang, this.device, program, inputsData, output, uniforms);
Expand Down
9 changes: 4 additions & 5 deletions tfjs-backend-webgpu/src/kernels/webgpu_program.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {DataType, Tensor} from '@tensorflow/tfjs-core';
import {DataType, Rank, ShapeMap, Tensor} from '@tensorflow/tfjs-core';
import {Glslang} from '@webgpu/glslang/dist/web-devel/glslang.onefile';

import * as shader_preprocessor from '../shader_preprocessor';
Expand Down Expand Up @@ -121,11 +121,10 @@ export const compileProgram =
return {bindGroupLayout, pipeline};
};

// TODO: Consider uploading shape info as vec4s regardless of rank to reduce
// recompilation.
export function makeShaderKey(program: WebGPUProgram, ranks: number[]): string {
export function makeShaderKey<R extends Rank>(program: WebGPUProgram,
shapes: Array<ShapeMap[R]>, types: string[]): string {
const key = (program.workGroupSize ? program.workGroupSize.join(',') : '') +
ranks.join(',') +
shapes.join(',') + types.join(',') + program.variableNames.join(',') +
(program.shaderKey ? program.shaderKey : program.userCode);
return key;
}