license import { GPGPUProgram } from './gpgpu_math' ; export class ResizeBilinearPackedProgram implements GPGPUProgram { variableNames = [ 'A' ] ; packedInputs = true ; packedOutput = true ; outputShape : number [ ] = [ ] ; userCode : string ; constructor ( inputShape , newHeight , newWidth , alignCorners , halfPixelCenters ) { const [ batch , oldHeight , oldWidth , depth ] = inputShape ; this . outputShape = [ batch , newHeight , newWidth , depth ] ; const effectiveInSize = [ ( alignCorners && newHeight > 1 ) ? oldHeight - 1 : oldHeight , ( alignCorners && newWidth > 1 ) ? oldWidth - 1 : oldWidth ] ; const effectiveOutSize = [ ( alignCorners && newHeight > 1 ) ? newHeight - 1 : newHeight , ( alignCorners && newWidth > 1 ) ? newWidth - 1 : newWidth ] ; let sourceFracIndexRC ; if ( halfPixelCenters ) { sourceFracIndexRC = `(vec3(yRC) + vec3(0.5)) * ` + `effectiveInputOverOutputRatioRC - vec3(0.5)` ; } else { sourceFracIndexRC = `vec3(yRC) * effectiveInputOverOutputRatioRC` ; } this . userCode = `
      const vec3 effectiveInputOverOutputRatioRC = vec3(
          ${ effectiveInSize [ 0 ] / effectiveOutSize [ 0 ] },
          ${ effectiveInSize [ 1 ] / effectiveOutSize [ 1 ] },
          ${ effectiveInSize [ 1 ] / effectiveOutSize [ 1 ] });
      const vec3 inputShapeRC = vec3(${ oldHeight }.0, ${ oldWidth }.0,
                                     ${ oldWidth }.0);

      float getAValue(int b, int r, int c, int d) {
        return getChannel(getA(b, r, c, d), vec2(c, d));
      }

      void main() {
        ivec4 coords = getOutputCoords();
        int b = coords[0];
        int d = coords[3];
        // Calculate values for next column in yRC.z.
        ivec3 yRC = coords.yzz + ivec3(0, 0, 1);

        // Fractional source index.
        vec3 sourceFracIndexRC = ${ sourceFracIndexRC };

        // Compute the four integer indices.
        ivec3 sourceFloorRC = ivec3(max(sourceFracIndexRC, vec3(0.0)));
        ivec3 sourceCeilRC = ivec3(
          min(inputShapeRC - 1.0, ceil(sourceFracIndexRC)));

        // Should we calculate next column and row elements in 2x2 packed cell.
        bool hasNextCol = d < ${ depth - 1 };
        bool hasNextRow = coords.z < ${ newWidth - 1 };

        // In parallel, construct four corners for all four components in
        // packed 2x2 cell.
        vec4 topLeft = vec4(
          getAValue(b, sourceFloorRC.x, sourceFloorRC.y, d),
          hasNextCol ? getAValue(b, sourceFloorRC.x, sourceFloorRC.y, d + 1)
                     : 0.0,
          hasNextRow ? getAValue(b, sourceFloorRC.x, sourceFloorRC.z, d)
                     : 0.0,
          (hasNextRow && hasNextCol) ?
            getAValue(b, sourceFloorRC.x, sourceFloorRC.z, d + 1) : 0.0);

        vec4 bottomLeft = vec4(
          getAValue(b, sourceCeilRC.x, sourceFloorRC.y, d),
          hasNextCol ? getAValue(b, sourceCeilRC.x, sourceFloorRC.y, d + 1)
                     : 0.0,
          hasNextRow ? getAValue(b, sourceCeilRC.x, sourceFloorRC.z, d)
                     : 0.0,
          (hasNextRow && hasNextCol) ?
            getAValue(b, sourceCeilRC.x, sourceFloorRC.z, d + 1) : 0.0);

        vec4 topRight = vec4(
          getAValue(b, sourceFloorRC.x, sourceCeilRC.y, d),
          hasNextCol ? getAValue(b, sourceFloorRC.x, sourceCeilRC.y, d + 1)
                     : 0.0,
          hasNextRow ? getAValue(b, sourceFloorRC.x, sourceCeilRC.z, d)
                     : 0.0,
          (hasNextRow && hasNextCol) ?
            getAValue(b, sourceFloorRC.x, sourceCeilRC.z, d + 1) : 0.0);

        vec4 bottomRight = vec4(
          getAValue(b, sourceCeilRC.x, sourceCeilRC.y, d),
          hasNextCol ? getAValue(b, sourceCeilRC.x, sourceCeilRC.y, d + 1)
                     : 0.0,
          hasNextRow ? getAValue(b, sourceCeilRC.x, sourceCeilRC.z, d)
                     : 0.0,
          (hasNextRow && hasNextCol) ?
            getAValue(b, sourceCeilRC.x, sourceCeilRC.z, d + 1) : 0.0);

        vec3 fracRC = sourceFracIndexRC - vec3(sourceFloorRC);

        vec4 top = mix(topLeft, topRight, fracRC.yyzz);
        vec4 bottom = mix(bottomLeft, bottomRight, fracRC.yyzz);
        vec4 newValue = mix(top, bottom, fracRC.x);

        setOutput(newValue);
      }
    ` ; } }