Function makeShader

tfjs-backend-webgpu/src/webgpu_program.ts:198–348

(
    inputInfo: InputInfo[], outputData: {dtype: DataType, shape: number[]},
    program: WebGPUProgram)

}

function makeShader(
    inputInfo: InputInfo[], outputData: {dtype: DataType, shape: number[]},
    program: WebGPUProgram): string {
  const prefixSnippets: string[] = [];
  const flatWorkgroupSize = program.workgroupSize[0] *
      program.workgroupSize[1] * program.workgroupSize[2];
  program.outputComponent =
      program.outputComponent ? program.outputComponent : 1;
  prefixSnippets.push(`

      var<private> localId: vec3<u32>;
      var<private> localIndex: u32;
      var<private> globalId: vec3<u32>;
      var<private> numWorkgroups: vec3<u32>;
      var<private> workgroupId: vec3<u32>;

      // Only used when the y/z dimension of workgroup size is 1.
      fn getGlobalIndex() -> i32 {
        ${
      isFlatDispatch(program) ?
          `  return i32(globalId.x);` :
          `  return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +
                workgroupId.y * numWorkgroups.x + workgroupId.x) * ${
              flatWorkgroupSize}u +
                localIndex);
        `}
      }
    `);

  if (program.pixelsOpType != null) {
    const inoutSnippet = program.pixelsOpType === PixelsOpType.FROM_PIXELS ?
        `@group(0) @binding(0) var<storage, read_write> result: array<${
            dataTypeToGPUType(outputData.dtype, program.outputComponent)}>;` :
        `@group(0) @binding(1) var<storage, read> inBuf : array<${
            dataTypeToGPUType(inputInfo[0].dtype, program.outputComponent)}>;`;
    const outShapeStridesType =
        outputData.shape.length === 3 ? 'vec2<i32>' : 'i32';
    prefixSnippets.push(`
        struct Uniform {
          outShapeStrides : ${outShapeStridesType},
          size            : i32,
          numChannels     : i32,
          alpha           : f32,
        };

        ${inoutSnippet}
        @group(0) @binding(2) var<uniform> uniforms: Uniform;
      `);
    const useGlobalIndex = isFlatDispatchLayout(program);
    return [
      commonSnippet,
      prefixSnippets.join('\n'),
      getCoordsFromIndexSnippet(outputData.shape),
      program.getUserCode(),
      getStartHeaderString(useGlobalIndex, program),
    ].join('\n');
  }