Function getOutputCoordsSnippet

tfjs-backend-webgpu/src/webgpu_program.ts:645–713

* Generates getOutputCoords() function that computes output coordinates * from dispatch geometry to reduce arithmetic.

(
    outShape: number[],
    dispatchLayout: {x: number[], y?: number[], z?: number[]})

 * from dispatch geometry to reduce arithmetic.
 */
function getOutputCoordsSnippet(
    outShape: number[],
    dispatchLayout: {x: number[], y?: number[], z?: number[]}): string {
  const {x, y = [], z = []} = dispatchLayout;

  const outRank = outShape.length;
  const rank = x.length + y.length + z.length;
  // getOutputCoords is only meaningful when the output rank is same with
  // dispatch layout rank.
  if (rank !== outRank) {
    return '';
  }

  if (x.length === outRank) {
    const dtype = getCoordsDataType(outRank);
    const snippet = `fn getOutputCoords() -> ${dtype}{
    let globalIndex = getGlobalIndex();
    return getCoordsFromIndex(globalIndex);
  }
  `;
    return snippet;
  }

  let gatherDimensionsStr = '';
  const dims = [x, y, z];

  for (let i = 0; i < dims.length; i++) {
    const arr = dims[i];

    if (arr.length === 0) {
      continue;
    }

    if (arr.length === 1) {
      gatherDimensionsStr += `let d${arr[0]} = i32(globalId[${i}]);`;
    } else {
      const strides = symbolicallyComputeStrides(arr, 'uniforms.outShape');
      gatherDimensionsStr += `var index${i} = i32(globalId[${i}]);`;
      for (let j = 0; j < strides.length; j++) {
        gatherDimensionsStr += `let d${arr[j]} = index${i} / ${strides[j]};`;

        if (j === strides.length - 1) {
          gatherDimensionsStr += `let d${arr[j + 1]} = ` +
              `index${i} - d${arr[j]} * ${strides[j]};`;
        } else {
          gatherDimensionsStr +=
              `index${i} = index${i} - d${arr[j]} * ${strides[j]};`;
        }
      }
    }
  }

  const dimensions = [];
  for (let i = 0; i < rank; i++) {
    dimensions.push(`d${i}`);
  }

  const dtype = getCoordsDataType(rank);