license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType , UniformType } from './shader_compiler' ; export class SliceProgram implements GPGPUProgram { variableNames = [ 'source' ] ; outputShape : number [ ] ; userCode : string ; rank : number ; customUniforms : Array < { name : string ; arrayIndex : number ; type : UniformType ; } > ; constructor ( destSize ) { this . outputShape = destSize ; this . rank = destSize . length ; const dtype = getCoordsDataType ( this . rank ) ; this . customUniforms = [ { name : 'start' , arrayIndex : this . rank , type : 'int' } ] ; const sourceCoords = getCoords ( this . rank ) ; let body ; const coordSum = destSize . map ( ( _ , i ) => { return `sourceLoc.${ coords [ i ] } = start[${ i }] + coords.${ coords [ i ] };` ; } ) ; body = `
        ${ dtype } sourceLoc;
        ${ dtype } coords = getOutputCoords();
        ${ coordSum . join ( '\n' ) }
      ` ; this . userCode = `
      void main() {
        ${ body }
        setOutput(getSource(${ sourceCoords }));
      }
    ` ; } } const coords = [ 'x' , 'y' , 'z' , 'w' , 'u' , 'v' ] ; function getCoords ( rank ) { if ( rank === 1 ) { return 'sourceLoc' ; } else if ( rank <= 6 ) { return coords . slice ( 0 , rank ) . map ( x => 'sourceLoc.' + x ) . join ( ',' ) ; } else { throw Error ( `Slicing for rank ${ rank } is not yet supported` ) ; } }