Skip to content
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

[webgpu] Tweak WGSL entry point emission so '{' isn't visually swallowed #6764

Merged
merged 3 commits into from
Aug 22, 2022
Merged

[webgpu] Tweak WGSL entry point emission so '{' isn't visually swallowed #6764

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f40116d
[webgpu] Tweak WGSL entry point emission so '{' isn't visually swallowed
hujiajie Aug 18, 2022
4020ced
Merge branch 'master' into master
Aug 19, 2022
5e893f0
Merge branch 'master' into master
qjia7 Aug 22, 2022
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
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/addn_packed_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class AddNPackedProgram implements WebGPUProgram {
Expand Down Expand Up @@ -52,7 +52,7 @@ export class AddNPackedProgram implements WebGPUProgram {
.join(' + ');

const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
for (var i = 0; i < ${this.workPerThread}; i = i + 1) {
let flatIndex = index * ${this.workPerThread} + i;
if (flatIndex < uniforms.size) {
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/argminmax_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
*/

import {backend_util} from '@tensorflow/tfjs-core';
import {getCoordsXYZ, getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getCoordsXYZ, getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class ArgMinMaxProgram implements WebGPUProgram {
Expand Down Expand Up @@ -91,7 +91,7 @@ export class ArgMinMaxProgram implements WebGPUProgram {

${sharedMemorySnippet}

${getMainHeaderAndGlobalIndexString()}
${main('index')} {
let outputIndex = index / i32(workGroupSizeX);
let reduceLength = ${getInputShapeLastDim()};

Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/batchnorm_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
*/

import {backend_util} from '@tensorflow/tfjs-core';
import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class BatchNormProgram implements WebGPUProgram {
Expand Down Expand Up @@ -69,7 +69,7 @@ export class BatchNormProgram implements WebGPUProgram {
}

const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size)
{
let xValue = getXByOutputIndex(index);
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/binary_op_complex_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

import {backend_util} from '@tensorflow/tfjs-core';
import {BinaryOpType, getBinaryOpString} from './binary_op_util';
import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class BinaryOpComplexProgram implements WebGPUProgram {
Expand Down Expand Up @@ -48,7 +48,7 @@ export class BinaryOpComplexProgram implements WebGPUProgram {
${opStr}
}

${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if(index < uniforms.size) {
let areal = getARealByOutputIndex(index);
let aimag = getAImagByOutputIndex(index);
Expand Down
7 changes: 3 additions & 4 deletions tfjs-backend-webgpu/src/binary_op_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import {backend_util, util} from '@tensorflow/tfjs-core';

import {BinaryOpType, getBinaryOpString} from './binary_op_util';
import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class BinaryOpProgram implements WebGPUProgram {
Expand Down Expand Up @@ -105,8 +105,7 @@ export class BinaryOpProgram implements WebGPUProgram {
${opStr}
}
var<workgroup> sharedBuf : array<f32, ${this.lastDimensionSize}>;
${getMainHeaderAndGlobalIndexString()}

${main('index')} {
// Fill in the shared memory buffer. Here we need a loop to make sure
// that all data in A|B are uploaded when |sharedMemorySize| is larger
// than work group size.
Expand Down Expand Up @@ -136,7 +135,7 @@ export class BinaryOpProgram implements WebGPUProgram {
fn binaryOperation(a : ${dType}, b : ${dType}) -> ${dType} {
${opStr}
}
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
let a = getAByOutputIndex(index);
let b = getBByOutputIndex(index);
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/clip_vec4_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class ClipVec4Program implements WebGPUProgram {
Expand All @@ -41,7 +41,7 @@ export class ClipVec4Program implements WebGPUProgram {

getUserCode(): string {
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if(index < uniforms.size) {
let value = getAByOutputIndex(index);
var clampedValue : vec4<f32>;
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/clip_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class ClipProgram implements WebGPUProgram {
Expand All @@ -41,7 +41,7 @@ export class ClipProgram implements WebGPUProgram {

getUserCode(): string {
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if(index < uniforms.size) {
let value = getAByOutputIndex(index);
if (isnan(value)) {
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/concat_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
*/

import {backend_util} from '@tensorflow/tfjs-core';
import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class ConcatProgram implements WebGPUProgram {
Expand Down Expand Up @@ -67,7 +67,7 @@ export class ConcatProgram implements WebGPUProgram {
}

const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
let flatIndex = index * ${this.workPerThread} + i;
if(flatIndex < uniforms.size) {
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/conv_backprop_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
*/

import {backend_util} from '@tensorflow/tfjs-core';
import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class Conv2DDerInputProgram implements WebGPUProgram {
Expand Down Expand Up @@ -45,7 +45,7 @@ export class Conv2DDerInputProgram implements WebGPUProgram {
const colDim = this.isChannelsLast ? 2 : 3;
const channelDim = this.isChannelsLast ? 3 : 1;
return `
${getMainHeaderAndGlobalIndexString()} {
${main('index')} {
if(index < uniforms.size) {
let coords = getCoordsFromIndex(index);
let batch = coords[0];
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/crop_and_resize_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class CropAndResizeProgram implements WebGPUProgram {
Expand Down Expand Up @@ -78,7 +78,7 @@ export class CropAndResizeProgram implements WebGPUProgram {
// tslint:disable-next-line:max-line-length
// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/crop_and_resize_op_gpu.cu.cc
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
let coords = getCoordsFromIndex(index);
let height_ratio = f32(${heightRatio});
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/cum_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export enum CumOpType {
Expand Down Expand Up @@ -70,7 +70,7 @@ export class CumProgram implements WebGPUProgram {
idxString = (this.reverse ? 'end + pow2' : 'end - pow2');
}
return `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
var coords = getCoordsFromIndex(index);

Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/depth_to_space_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class DepthToSpaceProgram implements WebGPUProgram {
Expand All @@ -40,7 +40,7 @@ export class DepthToSpaceProgram implements WebGPUProgram {

getUserCode(): string {
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
let coords = getCoordsFromIndex(index);
let b = coords[0];
Expand Down
8 changes: 4 additions & 4 deletions tfjs-backend-webgpu/src/depthwise_conv2d_nchw_shared_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,10 +83,10 @@ export class DepthwiseConv2DNCHWSharedProgram implements WebGPUProgram {
}

${getWorkGroupSizeString()}
fn main(@builtin(local_invocation_id) LocalId : vec3<u32>,
@builtin(global_invocation_id) GlobalId : vec3<u32>,
@builtin(local_invocation_index) LocalIndex: u32,
@builtin(num_workgroups) NumWorkgroups: vec3<u32>) {
fn _start(@builtin(local_invocation_id) LocalId : vec3<u32>,
@builtin(global_invocation_id) GlobalId : vec3<u32>,
@builtin(local_invocation_index) LocalIndex: u32,
@builtin(num_workgroups) NumWorkgroups: vec3<u32>) {
localId = LocalId;
globalId = GlobalId;
let localIndex = i32(LocalIndex);
Expand Down
2 changes: 1 addition & 1 deletion tfjs-backend-webgpu/src/depthwise_conv2d_vec4_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ export class DepthwiseConv2DVec4Program implements WebGPUProgram {
return value;
}
${getWorkGroupSizeString()}
fn main(@builtin(global_invocation_id) globalId: vec3<u32>) {
fn _start(@builtin(global_invocation_id) globalId: vec3<u32>) {
let batch = i32(globalId.z) / uniforms.outShape[1];
let r = i32(globalId.z) % uniforms.outShape[1];
let c = i32(globalId.y) * 4;
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/depthwise_conv2d_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import {backend_util} from '@tensorflow/tfjs-core';

import {activationFnSnippet, biasActivationSnippet} from './activation_util';
import {getMainHeaderString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class DepthwiseConv2DProgram implements WebGPUProgram {
Expand Down Expand Up @@ -67,7 +67,7 @@ export class DepthwiseConv2DProgram implements WebGPUProgram {
const userCode = `
${activationFnSnippet(this.activation, this.hasPreluActivation, false, 4)}

${getMainHeaderString()}
${main()} {
let coords = getOutputCoords();
let batch = coords[0];
let xRCCorner = vec2<i32>(coords.${
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/fill_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class FillProgram implements WebGPUProgram {
Expand All @@ -39,7 +39,7 @@ export class FillProgram implements WebGPUProgram {

getUserCode(): string {
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
setOutputAtIndex(index, uniforms.value);
}
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/flip_left_right_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class FlipLeftRightProgram implements WebGPUProgram {
Expand All @@ -37,7 +37,7 @@ export class FlipLeftRightProgram implements WebGPUProgram {

getUserCode(): string {
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
let coords = getCoordsFromIndex(index);
let coordX = uniforms.xShape[2] - coords[2] - 1;
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/from_pixels_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class FromPixelsProgram implements WebGPUProgram {
Expand Down Expand Up @@ -48,7 +48,7 @@ export class FromPixelsProgram implements WebGPUProgram {
this.importVideo ? 'texture_external' : 'texture_2d<f32>';
return `
@binding(1) @group(0) var src: ${textureType};
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
let flatIndex = index * uniforms.numChannels;
if (flatIndex < uniforms.size) {
let coords = getCoordsFromIndex(flatIndex);
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/gather_nd_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getCoordsDataType, getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getCoordsDataType, getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class GatherNDProgram implements WebGPUProgram {
Expand Down Expand Up @@ -46,7 +46,7 @@ export class GatherNDProgram implements WebGPUProgram {
strideString = 'uniforms.strides';
}
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
let coords = getCoordsFromIndex(index);
var flattenIndex = 0;
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-webgpu/src/gather_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
* =============================================================================
*/

import {getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class GatherProgram implements WebGPUProgram {
Expand All @@ -41,7 +41,7 @@ export class GatherProgram implements WebGPUProgram {
getUserCode(): string {
const sourceCoords = getSourceCoords(this.aShape);
const userCode = `
${getMainHeaderAndGlobalIndexString()}
${main('index')} {
if (index < uniforms.size) {
let resRC = getCoordsFromIndex(index);
let indexZ = i32(getIndices(resRC.x, resRC.z));
Expand Down
Loading