Method getUserCode

tfjs-backend-webgpu/src/scatter_webgpu.ts:61–124

()

  }

  getUserCode(): string {
    let indicesString = '';
    if (this.indicesRank === 1) {
      indicesString = 'coords[0]';
    } else if (this.indicesRank === 2) {
      indicesString = 'coords[0], j';
    }
    const indicesSnippet = `getIndices(${indicesString})`;

    const strideString = this.sliceDimGreaterThanOne ? 'uniforms.strides[j]' :
                                                       'uniforms.strides';

    let outCoordsString = '';
    let getUpdatesCoordsFromFlatIndex = '';
    if (this.dispatchLayout.x.length === 1) {
      outCoordsString = 'flattenedIndex';
      getUpdatesCoordsFromFlatIndex = `
      fn getUpdatesCoordsFromFlatIndex(index : i32) -> i32 {
        return index;
      }
      `;
    } else if (this.dispatchLayout.x.length === 2) {
      outCoordsString = 'vec2<i32>(flattenedIndex, coords[1])';
      getUpdatesCoordsFromFlatIndex = `
      fn getUpdatesCoordsFromFlatIndex(index : i32) -> vec2<i32> {
        // N.B. |updates| could be a scalar tensor, conceptually representing a
        // 2D tensor with all values equal to that. By design, its size must be
        // the same as |outShape[1]| in one dimension, and |indicesShape[0]|
        // gives the other.
        let sliceSize = uniforms.outShape[1];
        let d0 = index / sliceSize;
        let d1 = index - d0 * sliceSize;
        return vec2<i32>(d0, d1);
      }
      `;
    }
    const updatesString =
        Array.from({length: this.updatesRank}, (_, idx) => `coords[${idx}]`);
    const updatesSnippet = `getUpdates(${updatesString.join(', ')})`;

    const userCode = `
    ${getUpdatesCoordsFromFlatIndex}
      ${main('index')} {
        if (index < uniforms.updatesSize) {
          let coords = getUpdatesCoordsFromFlatIndex(index);
          var flattenedIndex = 0;
          for (var j = 0; j < uniforms.sliceDim; j = j + 1) {
            let indexInside = i32(round(${indicesSnippet}));
            flattenedIndex = flattenedIndex + indexInside * ${strideString};
          }
          let updateValue =
              ${dataTypeToGPUType(this.type)}(${updatesSnippet});
          let flatIndex = getOutputIndexFromCoords(${outCoordsString});

          ${
        this.sumDupeIndices ?
            atomicAddSnippet(
                '&result[flatIndex]', 'updateValue',