license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class SelectProgram implements GPGPUProgram { variableNames = [ 'c' , 'a' , 'b' ] ; outputShape : number [ ] ; userCode : string ; constructor ( cRank , shape , rank ) { this . outputShape = shape ; let cCoords ; let abCoords ; if ( rank > 4 ) { throw Error ( `Where for rank ${ rank } is not yet supported` ) ; } if ( rank === 1 ) { abCoords = `resRC` ; cCoords = `resRC` ; } else { const currentCoords = [ 'resRC.x' , 'resRC.y' , 'resRC.z' , 'resRC.w' ] ; const cCoordVars = [ ] ; const abCoordVars = [ ] ; for ( let i = 0 ; i < shape . length ; i ++ ) { abCoordVars . push ( `${ currentCoords [ i ] }` ) ; if ( i < cRank ) { cCoordVars . push ( `${ currentCoords [ i ] }` ) ; } } cCoords = cCoordVars . join ( ) ; abCoords = abCoordVars . join ( ) ; } const dtype = getCoordsDataType ( rank ) ; this . userCode = `
      void main() {
        ${ dtype } resRC = getOutputCoords();
        float cVal = getC(${ cCoords });
        if (cVal >= 1.0) {
          setOutput(getA(${ abCoords }));
        } else {
          setOutput(getB(${ abCoords }));
        }
      }
    ` ; } }