[webgpu] Tweak WGSL entry point emission so '{' isn't visually swallo…

…wed (#6764)
tensorflow · Aug 22, 2022 · 2a25492 · 2a25492
1 parent d86a6db
commit 2a25492
Show file tree

Hide file tree

Showing 43 changed files with 134 additions and 116 deletions.
diff --git a/tfjs-backend-webgpu/src/addn_packed_webgpu.ts b/tfjs-backend-webgpu/src/addn_packed_webgpu.ts
@@ -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 {
@@ -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) {

diff --git a/tfjs-backend-webgpu/src/argminmax_webgpu.ts b/tfjs-backend-webgpu/src/argminmax_webgpu.ts
@@ -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 {
@@ -91,7 +91,7 @@ export class ArgMinMaxProgram implements WebGPUProgram {
 
       ${sharedMemorySnippet}
 
-      ${getMainHeaderAndGlobalIndexString()}
+      ${main('index')} {
         let outputIndex = index / i32(workGroupSizeX);
         let reduceLength = ${getInputShapeLastDim()};
 

diff --git a/tfjs-backend-webgpu/src/batchnorm_webgpu.ts b/tfjs-backend-webgpu/src/batchnorm_webgpu.ts
@@ -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 {
@@ -69,7 +69,7 @@ export class BatchNormProgram implements WebGPUProgram {
     }
 
     const userCode = `
-      ${getMainHeaderAndGlobalIndexString()}
+      ${main('index')} {
         if (index < uniforms.size)
         {
           let xValue = getXByOutputIndex(index);

diff --git a/tfjs-backend-webgpu/src/binary_op_complex_webgpu.ts b/tfjs-backend-webgpu/src/binary_op_complex_webgpu.ts
@@ -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 {
@@ -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);

diff --git a/tfjs-backend-webgpu/src/binary_op_webgpu.ts b/tfjs-backend-webgpu/src/binary_op_webgpu.ts
@@ -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 {
@@ -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.
@@ -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);

diff --git a/tfjs-backend-webgpu/src/clip_vec4_webgpu.ts b/tfjs-backend-webgpu/src/clip_vec4_webgpu.ts
@@ -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 {
@@ -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>;

diff --git a/tfjs-backend-webgpu/src/clip_webgpu.ts b/tfjs-backend-webgpu/src/clip_webgpu.ts
@@ -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 {
@@ -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)) {

diff --git a/tfjs-backend-webgpu/src/concat_webgpu.ts b/tfjs-backend-webgpu/src/concat_webgpu.ts
@@ -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 {
@@ -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) {

diff --git a/tfjs-backend-webgpu/src/conv_backprop_webgpu.ts b/tfjs-backend-webgpu/src/conv_backprop_webgpu.ts
@@ -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 {
@@ -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];

diff --git a/tfjs-backend-webgpu/src/crop_and_resize_webgpu.ts b/tfjs-backend-webgpu/src/crop_and_resize_webgpu.ts
@@ -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 {
@@ -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});

diff --git a/tfjs-backend-webgpu/src/cum_webgpu.ts b/tfjs-backend-webgpu/src/cum_webgpu.ts
@@ -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 {
@@ -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);
 

diff --git a/tfjs-backend-webgpu/src/depth_to_space_webgpu.ts b/tfjs-backend-webgpu/src/depth_to_space_webgpu.ts
@@ -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 {
@@ -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];

diff --git a/tfjs-backend-webgpu/src/depthwise_conv2d_nchw_shared_webgpu.ts b/tfjs-backend-webgpu/src/depthwise_conv2d_nchw_shared_webgpu.ts
@@ -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);

diff --git a/tfjs-backend-webgpu/src/depthwise_conv2d_vec4_webgpu.ts b/tfjs-backend-webgpu/src/depthwise_conv2d_vec4_webgpu.ts
@@ -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;

diff --git a/tfjs-backend-webgpu/src/depthwise_conv2d_webgpu.ts b/tfjs-backend-webgpu/src/depthwise_conv2d_webgpu.ts
@@ -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 {
@@ -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.${

diff --git a/tfjs-backend-webgpu/src/fill_webgpu.ts b/tfjs-backend-webgpu/src/fill_webgpu.ts
@@ -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 {
@@ -39,7 +39,7 @@ export class FillProgram implements WebGPUProgram {
 
   getUserCode(): string {
     const userCode = `
-    ${getMainHeaderAndGlobalIndexString()}
+    ${main('index')} {
       if (index < uniforms.size) {
         setOutputAtIndex(index, uniforms.value);
       }

diff --git a/tfjs-backend-webgpu/src/flip_left_right_webgpu.ts b/tfjs-backend-webgpu/src/flip_left_right_webgpu.ts
@@ -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 {
@@ -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;

diff --git a/tfjs-backend-webgpu/src/from_pixels_webgpu.ts b/tfjs-backend-webgpu/src/from_pixels_webgpu.ts
@@ -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 {
@@ -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);

diff --git a/tfjs-backend-webgpu/src/gather_nd_webgpu.ts b/tfjs-backend-webgpu/src/gather_nd_webgpu.ts
@@ -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 {
@@ -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;

diff --git a/tfjs-backend-webgpu/src/gather_webgpu.ts b/tfjs-backend-webgpu/src/gather_webgpu.ts
@@ -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 {
@@ -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));