license import { backend_util , util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class MeanProgram implements GPGPUProgram { variableNames = [ 'x' ] ; outputShape : number [ ] ; userCode : string ; constructor ( reduceInfo , divisor ? ) { const { windowSize , batchSize , inSize , outSize } = reduceInfo ; this . outputShape = [ batchSize , outSize ] ; const windowSizeNearestVec4 = Math . floor ( windowSize / 4 ) * 4 ; const windowSizeVec4Remainder = windowSize % 4 ; let updateSnippet = `sumValue += dot(values, ones);` ; if ( divisor != null ) { const denominator = 1 / divisor ; updateSnippet = `sumValue += dot(values * ${ util . isInt ( denominator ) ? denominator . toPrecision ( 2 ) : denominator }, ones);` ; } let checkOutOfBounds = '' ; if ( inSize % windowSize > 0 ) { checkOutOfBounds = `
        if (inIdx < 0 || inIdx >= ${ inSize }) {
          return 0.0;
        }
      ` ; } this . userCode = `
      const vec4 ones = vec4(1.0, 1.0, 1.0, 1.0);

      float getValue(int batch, int inIdx) {
        ${ checkOutOfBounds }
        return getX(batch, inIdx);
      }

      void main() {
        ivec2 coords = getOutputCoords();
        int batch = coords[0];
        int outIdx = coords[1];
        int inOffset = outIdx * ${ windowSize };

        float sumValue = 0.0;

        for (int i = 0; i < ${ windowSizeNearestVec4 }; i += 4) {
          int inIdx = inOffset + i;
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2),
            getValue(batch, inIdx + 3)
          );

          ${ updateSnippet }
        }

        int inIdx = inOffset + ${ windowSizeNearestVec4 };
        if (${ windowSizeVec4Remainder === 1 }) {
          vec4 values = vec4(getValue(batch, inIdx), 0.0, 0.0, 0.0);

          ${ updateSnippet }
        } else if (${ windowSizeVec4Remainder === 2 }) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1), 0.0, 0.0);

          ${ updateSnippet }
        } else if (${ windowSizeVec4Remainder === 3 }) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2), 0.0);

          ${ updateSnippet }
        }
        setOutput(sumValue);
      }
    ` ; } }