[webgpu] Update shader to support non module-level scoping function

FIXES tensorflow#6842 To support shader translation library which does not implement module scoping like naga
haoyunfeix · Oct 12, 2022 · 1f5c878 · 1f5c878
1 parent 2cc528b
commit 1f5c878
Show file tree

Hide file tree

Showing 41 changed files with 117 additions and 51 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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class AddNPackedProgram implements WebGPUProgram {
@@ -62,6 +62,7 @@ export class AddNPackedProgram implements WebGPUProgram {
           }
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }

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, util} from '@tensorflow/tfjs-core';
-import {getCoordsXYZ, getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getCoordsXYZ, getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class ArgMinMaxProgram implements WebGPUProgram {
@@ -140,6 +140,7 @@ export class ArgMinMaxProgram implements WebGPUProgram {
           setOutputAtIndexI32(outputIndex, xBestIndices[localId.x]);
         }
       }
+      ${start('index')}
     `;
       return userCode;
     } else {
@@ -160,6 +161,7 @@ export class ArgMinMaxProgram implements WebGPUProgram {
           setOutputAtIndexI32(index, bestIndex);
         }
       }
+      ${start('index')}
       `;
       return userCode;
     }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class BatchNormProgram implements WebGPUProgram {
@@ -81,6 +81,7 @@ export class BatchNormProgram implements WebGPUProgram {
           setOutputAtIndex(index,dot(vec3<f32>(xValue, -meanValue, offsetValue), vec3<f32>(inv, inv, 1.0)));
         }
       }
+      ${start('index')}
   `;
     return userCode;
   }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class BinaryOpComplexProgram implements WebGPUProgram {
@@ -57,6 +57,7 @@ export class BinaryOpComplexProgram implements WebGPUProgram {
           setOutputAtIndex(index, binaryOpComplex(areal, aimag, breal, bimag));
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class BinaryOpProgram implements WebGPUProgram {
@@ -118,6 +118,7 @@ export class BinaryOpProgram implements WebGPUProgram {
             setOutputAtIndex(index, binaryOperation(a, b));
           }
         }
+        ${start('index')}
         `;
     } else {
       userCode = `
@@ -129,6 +130,7 @@ export class BinaryOpProgram implements WebGPUProgram {
            setOutputAtIndex(index, binaryOperation(a, b));
          }
        }
+       ${start('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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class ClipVec4Program implements WebGPUProgram {
@@ -56,6 +56,7 @@ export class ClipVec4Program implements WebGPUProgram {
           setOutputAtIndex(index, clampedValue);
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }

diff --git a/tfjs-backend-webgpu/src/clip_webgpu.ts b/tfjs-backend-webgpu/src/clip_webgpu.ts
@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class ClipProgram implements WebGPUProgram {
@@ -51,6 +51,7 @@ export class ClipProgram implements WebGPUProgram {
           setOutputAtIndex(index, clamp(value, uniforms.minVal, uniforms.maxVal));
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class ConcatProgram implements WebGPUProgram {
@@ -79,6 +79,7 @@ export class ConcatProgram implements WebGPUProgram {
           }
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }

diff --git a/tfjs-backend-webgpu/src/conv2d_naive_webgpu.ts b/tfjs-backend-webgpu/src/conv2d_naive_webgpu.ts
@@ -18,7 +18,7 @@
 import {backend_util} from '@tensorflow/tfjs-core';
 
 import {activationFnSnippet, biasActivationSnippet} from './activation_util';
-import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch} from './webgpu_util';
 
 export class Conv2DNaiveProgram implements WebGPUProgram {
@@ -115,6 +115,7 @@ export class Conv2DNaiveProgram implements WebGPUProgram {
          }
          writeResult(batch, outRow, outCol, outChannel, acc);
        }
+       ${start('index')}
      `;
     return userCode;
   }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class Conv2DDerInputProgram implements WebGPUProgram {
@@ -94,6 +94,7 @@ export class Conv2DDerInputProgram implements WebGPUProgram {
         setOutputAtIndex(index, dotProd);
       }
     }
+    ${start('index')}
   `;
   }
 }
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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class CropAndResizeProgram implements WebGPUProgram {
@@ -133,6 +133,7 @@ export class CropAndResizeProgram implements WebGPUProgram {
         }
       }
     }
+    ${start('index')}
     `;
     return userCode;
   }

diff --git a/tfjs-backend-webgpu/src/cum_webgpu.ts b/tfjs-backend-webgpu/src/cum_webgpu.ts
@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export enum CumOpType {
@@ -85,6 +85,7 @@ export class CumProgram implements WebGPUProgram {
          setOutputAtIndex(index, val);
        }
       }
+      ${start('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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class DepthToSpaceProgram implements WebGPUProgram {
@@ -59,7 +59,8 @@ export class DepthToSpaceProgram implements WebGPUProgram {
           let rlt = ${this.getInputSamplingString()};
           setOutputAtIndex(index, rlt);
         }
-      }`;
+      }
+      ${start('index')}`;
     return userCode;
   }
 

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 as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class DepthwiseConv2DProgram implements WebGPUProgram {
@@ -65,7 +65,9 @@ export class DepthwiseConv2DProgram implements WebGPUProgram {
                                               'getX(batch, d1, xR, xC);';
 
     const userCode = `
-      ${activationFnSnippet(this.activation, this.hasPreluActivation, false, 4)}
+      ${
+        activationFnSnippet(
+            this.activation, this.hasPreluActivation, false, 4)}
 
       ${main()} {
         let coords = getOutputCoords();
@@ -131,6 +133,7 @@ export class DepthwiseConv2DProgram implements WebGPUProgram {
           setOutputAtCoords(coords[0], coords[1], coords[2], coords[3], value);
         }
       }
+      ${start()}
     `;
     return userCode;
   }

diff --git a/tfjs-backend-webgpu/src/fill_webgpu.ts b/tfjs-backend-webgpu/src/fill_webgpu.ts
@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class FillProgram implements WebGPUProgram {
@@ -44,6 +44,7 @@ export class FillProgram implements WebGPUProgram {
         setOutputAtIndex(index, uniforms.value);
       }
     }
+    ${start('index')}
   `;
     return userCode;
   }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class FlipLeftRightProgram implements WebGPUProgram {
@@ -45,6 +45,7 @@ export class FlipLeftRightProgram implements WebGPUProgram {
           setOutputAtIndex(index, outputValue);
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }

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 {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class FromPixelsProgram implements WebGPUProgram {
@@ -58,6 +58,7 @@ export class FromPixelsProgram implements WebGPUProgram {
           }
         }
       }
+      ${start('index')}
   `;
   }
 }
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, getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getCoordsDataType, getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class GatherNDProgram implements WebGPUProgram {
@@ -59,6 +59,7 @@ export class GatherNDProgram implements WebGPUProgram {
           setOutputAtIndex(index, getA(flattenIndex, coords[1]));
         }
       }
+      ${start('index')}
       `;
     return userCode;
   }

diff --git a/tfjs-backend-webgpu/src/gather_webgpu.ts b/tfjs-backend-webgpu/src/gather_webgpu.ts
@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, getStartHeaderString as start, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class GatherProgram implements WebGPUProgram {
@@ -49,6 +49,7 @@ export class GatherProgram implements WebGPUProgram {
           setOutputAtIndex(index, inBounds * getA(${sourceCoords}));
         }
       }
+      ${start('index')}
     `;
     return userCode;
   }