Function makeShaderKey

tfjs-backend-webgpu/src/webgpu_program.ts:350–381

(
    program: WebGPUProgram, inputsData: InputInfo[],
    output: TensorInfo)

}

export function makeShaderKey<R extends Rank>(
    program: WebGPUProgram, inputsData: InputInfo[],
    output: TensorInfo): string {
  let key = program.shaderKey;
  if (program.pixelsOpType != null) {
    return key;
  }

  const shapes: number[][] = [];
  const types: Array<keyof DataTypeMap> = [];
  inputsData.forEach(element => {
    shapes.push(element.shape);
    types.push(element.dtype);
  });
  shapes.push(output.shape);
  types.push(output.dtype);

  const broadcastDims =
      inputsData.map(d => backend_util.getBroadcastDims(d.shape, output.shape));
  const inputShapesEqualsOutShape =
      inputsData.map(d => util.arraysEqual(d.shape, output.shape)).join('_');
  const broadcastDimsKey = broadcastDims.map(d => d.join('_')).join(';');

  const flatDispatchString = isFlatDispatch(program) ? 'flatDispatch' : '';

  key += '_' + (program.workgroupSize ? program.workgroupSize.join(',') : '') +
      shapes.map(shape => shape.length).join(',') + types.join(',') +
      program.variableNames.join(',') + broadcastDimsKey +
      inputShapesEqualsOutShape + flatDispatchString;

  return key;
}

const commonSnippet = `
  struct vec5 {x: i32, y: i32, z: i32, w: i32, u: i32};