license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class TransposeProgram implements GPGPUProgram { variableNames = [ 'A' ] ; outputShape : number [ ] ; userCode : string ; rank : number ; constructor ( aShape , newDim ) { const outputShape = new Array ( aShape . length ) ; for ( let i = 0 ; i < outputShape . length ; i ++ ) { outputShape [ i ] = aShape [ newDim [ i ] ] ; } this . outputShape = outputShape ; this . rank = outputShape . length ; const dtype = getCoordsDataType ( this . rank ) ; const switched = getSwitchedCoords ( newDim ) ; this . userCode = `
    void main() {
      ${ dtype } resRC = getOutputCoords();
      setOutput(getA(${ switched }));
    }
    ` ; } } function getSwitchedCoords ( newDim ) { const rank = newDim . length ; if ( rank > 6 ) { throw Error ( `Transpose for rank ${ rank } is not yet supported` ) ; } const originalOrder = [ 'resRC.x' , 'resRC.y' , 'resRC.z' , 'resRC.w' , 'resRC.u' , 'resRC.v' ] ; const switchedCoords = new Array ( rank ) ; for ( let i = 0 ; i < newDim . length ; i ++ ) { switchedCoords [ newDim [ i ] ] = originalOrder [ i ] ; } return switchedCoords . join ( ) ; }