license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class TileProgram implements GPGPUProgram { variableNames = [ 'A' ] ; outputShape : number [ ] ; userCode : string ; rank : number ; constructor ( aShape , reps ) { const outputShape = new Array ( aShape . length ) ; for ( let i = 0 ; i < outputShape . length ; i ++ ) { outputShape [ i ] = aShape [ i ] * reps [ i ] ; } this . outputShape = outputShape ; this . rank = outputShape . length ; const dtype = getCoordsDataType ( this . rank ) ; const sourceCoords = getSourceCoords ( aShape ) ; this . userCode = `
      void main() {
        ${ dtype } resRC = getOutputCoords();
        setOutput(getA(${ sourceCoords }));
      }
    ` ; } } function getSourceCoords ( aShape ) { const rank = aShape . length ; if ( rank > 5 ) { throw Error ( `Tile for rank ${ rank } is not yet supported` ) ; } if ( rank === 1 ) { return `imod(resRC, ${ aShape [ 0 ] })` ; } const currentCoords = [ 'resRC.x' , 'resRC.y' , 'resRC.z' , 'resRC.w' , 'resRC.u' ] ; const sourceCoords = [ ] ; for ( let i = 0 ; i < aShape . length ; i ++ ) { sourceCoords . push ( `imod(${ currentCoords [ i ] }, ${ aShape [ i ] })` ) ; } return sourceCoords . join ( ) ; }