Function setOutputSnippet

tfjs-backend-webgpu/src/webgpu_program.ts:791–827

(
    outShape: number[], outBufferType: DataType, component: number)

}

function setOutputSnippet(
    outShape: number[], outBufferType: DataType, component: number): string {
  const outRank = outShape.length;
  const gpuType = dataTypeToGPUType(outBufferType, component);
  let snippet =
      `fn setOutputAtIndex(flatIndex : i32, value : ${typeSnippet(component)}) {
      result[flatIndex] = ${gpuType}(value);
    }

    fn setOutputAtIndexI32(flatIndex : i32, value : ${
          typeSnippet(component, 'i32')}) {
      result[flatIndex] = ${gpuType}(value);
    }
    `;
  if (outRank >= 2) {
    const dims = ['d0', 'd1', 'd2', 'd3', 'd4', 'd5'].slice(0, outRank);
    const type = getCoordsDataType(outRank);

    snippet += `
      fn setOutputAtCoords(${dims.map(d => `${d} : i32`).join(', ')}, value : ${
        typeSnippet(component)}) {
        let flatIndex = getOutputIndexFromCoords(${type}(${dims.join(', ')}));
        setOutputAtIndex(flatIndex${
        component === 1 ? '' : ` / ${component}`}, value);
      }
      fn setOutputAtCoordsI32(${
        dims.map(d => `${d} : i32`).join(', ')}, value : ${
        typeSnippet(component, 'i32')}) {
        let flatIndex = getOutputIndexFromCoords(${type}(${dims.join(', ')}));
        setOutputAtIndexI32(flatIndex${
        component === 1 ? '' : ` / ${component}`}, value);
      }
    `;
  }

  return snippet;
}

function insertAlignment(uniformShader: string) {
  // insert alignment when current pattern is vec5 or vec6