Method getUserCode

tfjs-backend-webgpu/src/pool_webgpu.ts:60–129

()

  }

  getUserCode(): string {
    let updateSnippet: string;
    if (this.poolType === 'avg') {
      updateSnippet = `resultValue = resultValue + value; count = count + 1.0;`;
    } else if (this.computePositions) {
      const positionStr = this.flattenPositions ?
          (this.includeBatchIndex ?
               `((batch * uniforms.xShape[1] + xR) * uniforms.xShape[2] + xC) * uniforms.xShape[3] + d` :
               `(xR * uniforms.xShape[2] + xC) * uniforms.xShape[3] + d`) :
          `wR * uniforms.filterDims.y + wC`;
      updateSnippet = `let currMaxValue = mix(value, maxValue, maxValueFound);
      if (value >= currMaxValue) {
        maxValue = value;
        maxValueFound = 1.0;
        maxPosition = ${positionStr};
      }`;
    } else {
      updateSnippet = `resultValue = max(value, resultValue);`;
    }

    let returnValue = `resultValue`;
    if (this.poolType === 'avg') {
      returnValue = `resultValue / max(count, 1.0)`;
    }

    const userCode = `
      ${main('index')} {
      if (index < uniforms.size) {
        let coords = getCoordsFromIndex(index);
          let batch = coords[0];
          let d = coords[3];
          let xRCCorner = vec2<i32>(coords.yz) * uniforms.strides - uniforms.pads;
          let xRCorner = xRCCorner.x;
          let xCCorner = xRCCorner.y;

          ${
        this.computePositions ?
            `var maxValue = 0.0;
            var maxValueFound = 0.0;
            var maxPosition = 0;` :
            `var resultValue = ${
                this.poolType === 'avg' ? '0.0' : '-1.0 / pow(10.0, -20.0)'};`}

          var count = 0.0;
          for (var wR = 0; wR < uniforms.filterDims.x; wR = wR + uniforms.dilations.x) {
            let xR = xRCorner + wR;

            if (xR < 0 || xR >= uniforms.convDims.x) {
              continue;
            }

            for (var wC = 0; wC < uniforms.filterDims.y; wC = wC + uniforms.dilations.y) {
              let xC = xCCorner + wC;
              if (xC < 0 || xC >= uniforms.convDims.y) {
                continue;
              }

              let value = getX(batch, xR, xC, d);