license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export type GatherShape = [ number , number , number , number ] ; export class GatherProgram implements GPGPUProgram { variableNames = [ 'A' , 'indices' ] ; outputShape : number [ ] ; userCode : string ; rank : number ; constructor ( aShape , outputShape ) { this . outputShape = outputShape ; this . rank = outputShape . length ; const dtype = getCoordsDataType ( this . rank ) ; const sourceCoords = getSourceCoords ( aShape , 2 ) ; this . userCode = `
      void main() {
        ${ dtype } resRC = getOutputCoords();
        int index = int(getIndices(resRC.x, resRC.z));
        float inBounds = (index >= 0) && (index < ${ aShape [ 2 ] }) ? 1.0 : 0.0;
        setOutput(inBounds * getA(${ sourceCoords }));
      }
    ` ; } } function getSourceCoords ( aShape , axis ) { const currentCoords = [ 'resRC.x' , 'resRC.y' , 'resRC.z' , 'resRC.w' ] ; const sourceCoords = [ ] ; for ( let i = 0 ; i < aShape . length ; i ++ ) { if ( i === 2 ) { sourceCoords . push ( 'index' ) ; } else { sourceCoords . push ( `${ currentCoords [ i ] }` ) ; } } return sourceCoords . join ( ) ; }