license import { GPGPUProgram , useShapeUniforms } from './gpgpu_math' ; import { getChannels , getSourceCoords } from './packing_util' ; import { getCoordsDataType } from './shader_compiler' ; export class UnpackProgram implements GPGPUProgram { variableNames = [ 'A' ] ; packedInputs = true ; packedOutput = false ; outputShape : number [ ] ; userCode : string ; enableShapeUniforms : boolean ; constructor ( outputShape ) { this . outputShape = outputShape ; this . enableShapeUniforms = useShapeUniforms ( this . outputShape . length ) ; const rank = outputShape . length ; const channels = getChannels ( 'rc' , rank ) ; const dtype = getCoordsDataType ( rank ) ; const sourceCoords = getSourceCoords ( rank , channels ) ; const innerDims = channels . slice ( - 2 ) ; const coords = rank <= 1 ? 'rc' : `vec2(${ innerDims . join ( ',' ) })` ; this . userCode = `
      void main() {
        ${ dtype } rc = getOutputCoords();
        vec4 packedInput = getA(${ sourceCoords });

        setOutput(getChannel(packedInput, ${ coords }));
      }
    ` ; } }