Method getUserCode

tfjs-backend-webgpu/src/slice_webgpu.ts:47–73

()

  }

  getUserCode(): string {
    const dtype = getCoordsDataType(this.rank);
    const sourceCoords = getCoords(this.rank);
    let coordSum;
    if (this.start.length === 1) {
      coordSum = this.outputShape.map((_, i) => {
        return `sourceLoc = uniforms.start + coords;`;
      });
    } else {
      coordSum = this.outputShape.map((_, i) => {
        return `sourceLoc.${coords[i]} = uniforms.start.${
            getCoordsXYZ(i)} + coords.${coords[i]};`;
      });
    }

    const userCode = `
      ${main('index')} {
        if (index < uniforms.size) {
          var sourceLoc : ${dtype};
          let coords = getCoordsFromIndex(index);
          ${coordSum.join('\n')}
          setOutputAtIndex(index, getSource(${sourceCoords}));
        }
      }
    `;
    return userCode;
  }
}

const coords = ['x', 'y', 'z', 'w', 'u', 'v'];