license import { GPGPUProgram } from './gpgpu_math' ; import { getChannels } from './packing_util' ; import { getCoordsDataType } from './shader_compiler' ; export class ReversePackedProgram implements GPGPUProgram { variableNames = [ 'x' ] ; outputShape : number [ ] ; userCode : string ; packedInputs = true ; packedOutput = true ; constructor ( xShape , axis ) { const rank = xShape . length ; if ( rank > 4 ) { throw new Error ( `WebGL backend: Reverse of rank-${ rank } tensor is not yet supported` ) ; } this . outputShape = xShape ; const channels = getChannels ( 'rc' , rank ) ; const nextColumn = `${ channels [ rank - 1 ] } + 1 < ${ this . outputShape [ rank - 1 ] }` ; const nextRow = `${ channels [ rank - 2 ] } + 1 < ${ this . outputShape [ rank - 2 ] }` ; const type = getCoordsDataType ( rank ) ; if ( rank === 1 ) { this . userCode = `
        void main(){
          int rc = getOutputCoords();
          vec4 result = vec4(0.);
          result.r = getChannel(getX(${ xShape [ 0 ] } - rc - 1),
            ${ xShape [ 0 ] } - rc - 1);
          if(${ nextColumn }){
              result.g = getChannel(getX(${ xShape [ 0 ] } - (rc  + 1) - 1),
                ${ xShape [ 0 ] } - (rc  + 1) - 1);
          }
          setOutput(result);
        }
      ` ; } else { this . userCode = `
        void main() {
          ${ type } rc = getOutputCoords();
          vec4 result = vec4(0.);
          result.r = ${ getR ( channels . slice ( ) ) };
          if(${ nextColumn }){
            result.g = ${ getG ( channels . slice ( ) ) };
          }
          if(${ nextRow }) {
            result.b = ${ getB ( channels . slice ( ) ) };
            if(${ nextColumn }) {
              result.a = ${ getA ( channels . slice ( ) ) };
            }
          }
          setOutput(result);
        }
    ` ; } function getR ( channels ) { return getChannel ( channels ) ; } function getG ( channels ) { channels [ rank - 1 ] = '(' + channels [ rank - 1 ] + ` + 1)` ; return getChannel ( channels ) ; } function getB ( channels ) { channels [ rank - 2 ] = '(' + channels [ rank - 2 ] + ` + 1)` ; return getChannel ( channels ) ; } function getA ( channels ) { channels [ rank - 1 ] = '(' + channels [ rank - 1 ] + ` + 1)` ; channels [ rank - 2 ] = '(' + channels [ rank - 2 ] + ` + 1)` ; return getChannel ( channels ) ; } function getChannel ( channels ) { const inCoordsArray = xShape . map ( ( _ , i ) => getInCoord ( i , channels ) ) ; const inCoords = inCoordsArray . join ( ',' ) ; const innerDims = inCoordsArray . slice ( - 2 ) . join ( ',' ) ; return `getChannel(getX(${ inCoords }), vec2(${ innerDims }))` ; } function getInCoord ( i , channels1 ) { if ( axis . indexOf ( i ) !== - 1 && xShape [ i ] !== 1 ) { return `${ xShape [ i ] } - ${ channels1 [ i ] } - 1` ; } else { return `${ channels1 [ i ] }` ; } } } }