license import { GPGPUProgram } from './gpgpu_math' ; import { getChannels } from './packing_util' ; import { getCoordsDataType , UniformType } from './shader_compiler' ; export class PadPackedProgram implements GPGPUProgram { variableNames = [ 'x' ] ; packedInputs = true ; packedOutput = true ; outputShape : number [ ] ; userCode : string ; customUniforms = [ { name : 'value' , type : 'float' as UniformType } ] ; constructor ( xShape , paddings , constantValue ) { this . outputShape = paddings . map ( ( p , i ) => p [ 0 ] + xShape [ i ] + p [ 1 ] ) ; const rank = xShape . length ; const dtype = getCoordsDataType ( rank ) ; const start = paddings . map ( p => p [ 0 ] ) . join ( ',' ) ; const end = paddings . map ( ( p , i ) => p [ 0 ] + xShape [ i ] ) . join ( ',' ) ; const coords = getChannels ( 'rc' , rank ) ; const source = getChannels ( 'source' , rank ) ; const cLimit = `${ coords [ rank - 1 ] } < ${ this . outputShape [ rank - 1 ] }` ; const innerDims = rank === 1 ? 'source' : `vec2(${ source . slice ( - 2 ) . join ( ) })` ; const componentSetup = [ `${ dtype } rc = outputLoc;` , `${ coords [ rank - 1 ] } += 1;
       if(${ cLimit }) {
      ` , rank === 1 ? '' : `}
       rc = outputLoc;
       ${ coords [ rank - 2 ] } += 1;
       if(${ coords [ rank - 2 ] } < ${ this . outputShape [ rank - 2 ] }) {` , rank === 1 ? '' : `  ${ coords [ rank - 1 ] } += 1;
         if(${ cLimit }) {` ] ; const paddingArea = rank === 1 ? 'rc < start || rc >= end' : 'any(lessThan(rc, start)) || any(greaterThanEqual(rc, end))' ; let mainLoop = '' ; for ( let i = 0 , j = rank === 1 ? 2 : 4 ; i < j ; i ++ ) { mainLoop += `
        ${ componentSetup [ i ] }
        if (${ paddingArea }) {
          result[${ i }] = float(value);
        } else {
          ${ dtype } source = rc - start;
          result[${ i }] = getChannel(getX(${ source . join ( ) }), ${ innerDims });
        }
      ` ; } mainLoop += ( rank === 1 ? `} ` : `}}` ) ; this . userCode = `
      const ${ dtype } start = ${ dtype }(${ start });
      const ${ dtype } end = ${ dtype }(${ end });

      void main() {
        ${ dtype } outputLoc = getOutputCoords();
        vec4 result = vec4(0.);
        ${ mainLoop }
        setOutput(result);
      }
    ` ; } }