license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class DepthwiseConv2DDerFilterProgram implements GPGPUProgram { variableNames = [ 'x' , 'dy' ] ; outputShape : number [ ] ; userCode : string ; constructor ( convInfo ) { this . outputShape = convInfo . filterShape ; const strideHeight = convInfo . strideHeight ; const strideWidth = convInfo . strideWidth ; const padTop = convInfo . padInfo . top ; const padLeft = convInfo . padInfo . left ; const channelMul = convInfo . outChannels / convInfo . inChannels ; this . userCode = `
      void main() {
        ivec4 coords = getOutputCoords();
        int wR = coords.x;
        int wC = coords.y;
        int d1 = coords.z;
        int dm = coords.w;
        int d2 = d1 * ${ channelMul } + dm;

        float dotProd = 0.0;

        // TO DO: Vec4 over the batch size
        for (int b = 0; b < ${ convInfo . batchSize }; b++) {
          for (int yR = 0; yR < ${ convInfo . outHeight }; yR++) {
            int xR = wR + yR * ${ strideHeight } - ${ padTop };

            if (xR < 0 || xR >= ${ convInfo . inHeight }) {
              continue;
            }

            for (int yC = 0; yC < ${ convInfo . outWidth }; yC++) {
              int xC = wC + yC * ${ strideWidth } - ${ padLeft };

              if (xC < 0 || xC >= ${ convInfo . inWidth }) {
                continue;
              }

              float dyValue = getDy(b, yR, yC, d2);
              float xValue = getX(b, xR, xC, d1);
              dotProd += (xValue * dyValue);
            }
          }
        }
        setOutput(dotProd);
      }
    ` ; } } export class DepthwiseConv2DDerInputProgram implements GPGPUProgram { variableNames = [ 'dy' , 'W' ] ; outputShape : number [ ] ; userCode : string ; constructor ( convInfo ) { this . outputShape = convInfo . inShape ; const filterHeight = convInfo . filterHeight ; const filterWidth = convInfo . filterWidth ; const strideHeight = convInfo . strideHeight ; const strideWidth = convInfo . strideWidth ; const padTop = filterHeight - 1 - convInfo . padInfo . top ; const padLeft = filterWidth - 1 - convInfo . padInfo . left ; const channelMul = convInfo . outChannels / convInfo . inChannels ; this . userCode = `
      const ivec2 pads = ivec2(${ padTop }, ${ padLeft });

      void main() {
        ivec4 coords = getOutputCoords();
        int batch = coords[0];
        int d1 = coords[3];
        ivec2 dyCorner = coords.yz - pads;
        int dyRCorner = dyCorner.x;
        int dyCCorner = dyCorner.y;

        float dotProd = 0.0;

        for (int wR = 0; wR < ${ filterHeight }; wR++) {
          float dyR = float(dyRCorner + wR) / ${ strideHeight }.0;

          if (dyR < 0.0 || dyR >= ${ convInfo . outHeight }.0 || fract(dyR) > 0.0) {
            continue;
          }
          int idyR = int(dyR);

          int wRPerm = ${ filterHeight } - 1 - wR;

          for (int wC = 0; wC < ${ filterWidth }; wC++) {
            float dyC = float(dyCCorner + wC) / ${ strideWidth }.0;

            if (dyC < 0.0 || dyC >= ${ convInfo . outWidth }.0 ||
                fract(dyC) > 0.0) {
              continue;
            }
            int idyC = int(dyC);

            int wCPerm = ${ filterWidth } - 1 - wC;

            // TO DO: Vec4 over the channelMul
            for (int dm = 0; dm < ${ channelMul }; dm++) {
              int d2 = d1 * ${ channelMul } + dm;
              float xValue = getDy(batch, idyR, idyC, d2);
              float wValue = getW(wRPerm, wCPerm, d1, dm);
              dotProd += xValue * wValue;
            }
          }
        }
        setOutput(dotProd);
      }
    ` ; } }