license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; import { getChannels } from './packing_util' ; import { getCoordsDataType } from './shader_compiler' ; export class ConcatPackedProgram implements GPGPUProgram { variableNames : string [ ] ; packedInputs = true ; packedOutput = true ; outputShape : number [ ] = [ ] ; userCode : string ; constructor ( shapes , axis ) { this . outputShape = backend_util . computeOutShape ( shapes , axis ) ; const shape = this . outputShape ; const rank = shape . length ; const dtype = getCoordsDataType ( rank ) ; const coords = getChannels ( 'coords' , rank ) ; const channels = [ 'x' , 'y' , 'z' , 'w' , 'u' , 'v' ] . slice ( 0 , rank ) ; this . variableNames = shapes . map ( ( _ , i ) => `T${ i }` ) ; const offsets = new Array ( shapes . length - 1 ) ; offsets [ 0 ] = shapes [ 0 ] [ axis ] ; for ( let i = 1 ; i < offsets . length ; i ++ ) { offsets [ i ] = offsets [ i - 1 ] + shapes [ i ] [ axis ] ; } const channel = channels [ axis ] ; const lastChannels = channels . slice ( - 2 ) ; const allChannels = channels . join ( ) ; let getValueSnippet = `if (${ channel } < ${ offsets [ 0 ] }) {
        return getChannel(
            getT0(${ allChannels }), vec2(${ lastChannels . join ( ) }));
        }` ; for ( let i = 1 ; i < offsets . length ; i ++ ) { const shift = offsets [ i - 1 ] ; getValueSnippet += `
        if (${ channel } < ${ offsets [ i ] }  && ${ channel } >= ${ offsets [ i - 1 ] }) {
          return getChannel(
            getT${ i }(${ shiftedChannels ( channels , channel , shift ) }),
            vec2(${ shiftedChannels ( lastChannels , channel , shift ) }));
        }` ; } const lastIndex = offsets . length ; const shift = offsets [ offsets . length - 1 ] ; getValueSnippet += `
        return getChannel(
          getT${ lastIndex }(${ shiftedChannels ( channels , channel , shift ) }),
          vec2(${ shiftedChannels ( lastChannels , channel , shift ) }));` ; this . userCode = `
      float getValue(${ channels . map ( x => 'int ' + x ) }) {
        ${ getValueSnippet }
      }

      void main() {
        ${ dtype } coords = getOutputCoords();
        vec4 result = vec4(getValue(${ coords }), 0., 0., 0.);

        ${ coords [ rank - 1 ] } = ${ coords [ rank - 1 ] } + 1;
        if (${ coords [ rank - 1 ] } < ${ shape [ rank - 1 ] }) {
          result.g = getValue(${ coords });
        }

        ${ coords [ rank - 2 ] } = ${ coords [ rank - 2 ] } + 1;
        if (${ coords [ rank - 2 ] } < ${ shape [ rank - 2 ] }) {
          result.a = getValue(${ coords });
        }

        ${ coords [ rank - 1 ] } = ${ coords [ rank - 1 ] } - 1;
        if (${ coords [ rank - 2 ] } < ${ shape [ rank - 2 ] } &&
            ${ coords [ rank - 1 ] } < ${ shape [ rank - 1 ] }) {
          result.b = getValue(${ coords });
        }
        setOutput(result);
      }
    ` ; } } param channels param channel param shift returns function shiftedChannels ( channels , channel , shift ) { const channelIdx = channels . indexOf ( channel ) ; const res = channels . map ( ( c , idx ) => { if ( idx === channelIdx ) { return `${ c } - ${ shift }` ; } else { return c ; } } ) ; return res . join ( ) ; }