license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class Dilation2DProgram implements GPGPUProgram { variableNames = [ 'x' , 'W' ] ; outputShape : number [ ] ; userCode : string ; constructor ( convInfo ) { this . outputShape = convInfo . outShape ; const { inHeight , inWidth , padInfo , strideHeight , strideWidth , filterHeight , filterWidth , dilationHeight , dilationWidth } = convInfo ; const { top : padTop , left : padLeft } = padInfo ; this . userCode = `
      const ivec2 strides = ivec2(${ strideHeight }, ${ strideWidth });
      const ivec2 pads = ivec2(${ padTop }, ${ padLeft });
      const float neg_infinity = -3.4e38;

      void main() {
        ivec4 coords = getOutputCoords();
        int batch = coords.x;
        int d1 = coords.w;
        ivec2 outTopLeftCorner =
            coords.yz * strides - pads;
        int hBeg = outTopLeftCorner.x;
        int wBeg = outTopLeftCorner.y;

        float curVal = neg_infinity;
        for (int h = 0; h < ${ filterHeight }; h++) {
          int hIn = hBeg + h * ${ dilationHeight };

          if (hIn >= 0 && hIn < ${ inHeight }) {
            for (int w = 0; w < ${ filterWidth }; w++) {
              int wIn = wBeg + w * ${ dilationWidth };

              if (wIn >= 0 && wIn < ${ inWidth }) {
                float xVal = getX(batch, hIn, wIn, d1);
                float wVal = getW(h, w, d1);

                float val = xVal + wVal;
                if (val > curVal) {
                  curVal = val;
                }
              }
            }
          }
        }

        float result = curVal;
        setOutput(result);
      }
    ` ; } }