license import { getGlslDifferences } from './glsl_version' ; import { GPGPUProgram , useShapeUniforms } from './gpgpu_math' ; import * as shader_util from './shader_compiler_util' ; export class EncodeMatrixPackedProgram implements GPGPUProgram { variableNames = [ 'A' ] ; userCode : string ; outputShape : number [ ] ; packedInputs = false ; packedOutput = true ; enableShapeUniforms : boolean ; customUniforms = [ { name : 'texShape' , type : 'ivec2' as const } ] ; constructor ( outputShape , inputIsUnsignedByte = false ) { const glsl = getGlslDifferences ( ) ; this . outputShape = outputShape ; this . enableShapeUniforms = useShapeUniforms ( this . outputShape . length ) ; let mainLoop = '' ; let output = 'result' ; if ( inputIsUnsignedByte ) { output = 'floor(result * 255. + 0.5)' ; } for ( let row = 0 ; row <= 1 ; row ++ ) { for ( let col = 0 ; col <= 1 ; col ++ ) { const channel = row * 2 + col ; mainLoop += `
          localCoords = coords;
          if(localCoords[2] + ${ col } < ${ this . enableShapeUniforms ? 'outShape[2]' : `${ outputShape [ 2 ] }` }) {
          localCoords[2] += ${ col };
          if (localCoords[1] + ${ row } < ${ this . enableShapeUniforms ? 'outShape[1]' : `${ outputShape [ 1 ] }` }) {
            localCoords[1] += ${ row };

            flatIndex = getFlatIndex(localCoords);
            offset = imod(flatIndex, 4);

            flatIndex = idiv(flatIndex, 4, 1.);

            int r = flatIndex / texShape[1];
            int c = imod(flatIndex, texShape[1]);
            vec2 uv = (vec2(c, r) + halfCR) / vec2(texShape[1], texShape[0]);
            values = ${ glsl . texture2D }(A, uv);

            if (offset == 0) {
              result[${ channel }] = values[0];
            } else if (offset == 1) {
              result[${ channel }] = values[1];
            } else if (offset == 2) {
              result[${ channel }] = values[2];
            } else {
              result[${ channel }] = values[3];
            }
          }
        }
        ` ; } } this . userCode = `
        ${ this . enableShapeUniforms ? shader_util . getFlatIndexFrom3DOutput ( ) : shader_util . getFlatIndexFrom3D ( outputShape ) }

        void main() {
          ivec3 coords = getOutputCoords();

          vec4 result = vec4(0.);
          int flatIndex, r, c, offset;
          ivec3 localCoords;
          vec2 uv;
          vec4 values;

          ${ mainLoop }

          ${ glsl . output } = ${ output };
        }
    ` ; } }