license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType , UniformType } from './shader_compiler' ; export class PadProgram implements GPGPUProgram { variableNames = [ 'x' ] ; outputShape : number [ ] ; userCode : string ; customUniforms = [ { name : 'value' , type : 'float' as UniformType } ] ; constructor ( xShape , paddings , constantValue ) { this . outputShape = paddings . map ( ( p , i ) => p [ 0 ] + xShape [ i ] + p [ 1 ] ) ; const rank = xShape . length ; const type = getCoordsDataType ( rank ) ; const start = paddings . map ( p => p [ 0 ] ) . join ( ',' ) ; const end = paddings . map ( ( p , i ) => p [ 0 ] + xShape [ i ] ) . join ( ',' ) ; const unpackedCoords = [ 'coords[0]' , 'coords[1]' , 'coords[2]' , 'coords[3]' ] . slice ( 0 , rank ) ; if ( rank === 1 ) { this . userCode = `
        int start = ${ start };
        int end = ${ end };

        void main() {
          int outC = getOutputCoords();
          if (outC < start || outC >= end) {
            setOutput(value);
          } else {
            setOutput(getX(outC - start));
          }
        }
      ` ; return ; } this . userCode = `
      ${ type } start = ${ type }(${ start });
      ${ type } end = ${ type }(${ end });

      void main() {
        ${ type } outC = getOutputCoords();
        if (any(lessThan(outC, start)) || any(greaterThanEqual(outC, end))) {
          setOutput(value);
        } else {
          ${ type } coords = outC - start;
          setOutput(getX(${ unpackedCoords }));
        }
      }
    ` ; } }