license import { GPGPUProgram , useShapeUniforms } from './gpgpu_math' ; import * as shader_util from './shader_compiler_util' ; export class ReshapePackedProgram implements GPGPUProgram { variableNames = [ 'A' ] ; packedInputs = true ; packedOutput = true ; outputShape : number [ ] ; userCode : string ; enableShapeUniforms : boolean ; customUniforms = [ { name : 'inputShape' , type : 'ivec3' as const } ] ; constructor ( outputShape , inputShape ) { this . outputShape = outputShape ; this . enableShapeUniforms = useShapeUniforms ( this . outputShape . length ) ; let mainLoop = `` ; for ( let i = 0 ; i < 4 ; i ++ ) { let thisRC = `thisRC = rc;` ; if ( i % 2 === 1 ) { thisRC += `thisRC.z += 1;` ; } if ( i > 1 ) { thisRC += `thisRC.y += 1;` ; } mainLoop += `
        ${ thisRC }
        ${ i > 0 ? `if(thisRC.y < rows && thisRC.z < cols){` : '' }
          int flatIndex = getFlatIndex(thisRC);

          ivec3 inputRC = inputCoordsFromReshapedOutCoords(flatIndex);
          vec2 inputRCInnerDims = vec2(float(inputRC.y),float(inputRC.z));

          result[${ i }] =
            getChannel(getA(inputRC.x, inputRC.y, inputRC.z), inputRCInnerDims);
        ${ i > 0 ? '}' : '' }
      ` ; } this . userCode = `
      ${ getReshapedInputCoords ( inputShape , this . enableShapeUniforms ) }
      ${ this . enableShapeUniforms ? shader_util . getFlatIndexFrom3DOutput ( ) : shader_util . getFlatIndexFrom3D ( outputShape ) }

      void main() {
        ivec3 rc = getOutputCoords();

        vec4 result = vec4(0.);

        ivec3 thisRC;
        int rows = ${ this . enableShapeUniforms ? 'outShape[1]' : outputShape [ 1 ] };
        int cols = ${ this . enableShapeUniforms ? 'outShape[2]' : outputShape [ 2 ] };

        ${ mainLoop }

        setOutput(result);
      }
    ` ; } } function getReshapedInputCoords ( shape , enableShapeUniforms ) { const coordsFromIndexSnippet = enableShapeUniforms ? shader_util . getLogicalCoordinatesFromFlatIndexByUniform ( [ 'r' , 'c' , 'd' ] , 'inputShape' ) : shader_util . getLogicalCoordinatesFromFlatIndex ( [ 'r' , 'c' , 'd' ] , shape ) ; return `
    ivec3 inputCoordsFromReshapedOutCoords(int index) {
      ${ coordsFromIndexSnippet }
      return ivec3(r, c, d);
    }
  ` ; }