[WebGL] Fix NHWC packed depthwise conv2d for dilation=3 (#6662)

FIX

* Update conv_packed_gpu_depthwise.ts

* add test

Author: Linchenn

tfjs-backend-webgl/src/conv_packed_gpu_depthwise.ts

@@ -193,18 +193,21 @@ export class DepthwiseConvPacked2DProgram implements GPGPUProgram {
               if (dilationWidth > 1) {
                 mainLoop += `
                     xCOffset -= 2;
-                    if (xCOffset >= 0 && xCOffset < inDims[1] && xTexelC${
-                    colIndex}Ready == 0) {
-                      xTexelC${colIndex} = getX(batch, xR, xCOffset, d1);
-                      xTexelC${colIndex}Ready = 1;
+                    if (xCOffset >= 0 && xCOffset < inDims[1]) {
+                     previous = getX(batch, xR, xCOffset, d1);
+                     xC${colIndex + 1} = vec4(previous.zw, xTexelC${
+                       colIndex + 1}.xy);
+                    } else {
+                     xC${colIndex + 1} = vec4(0.0, 0.0, xTexelC${
+                       colIndex + 1}.xy);
                     }
                     `;
+              } else {
+                mainLoop += `
+                    xC${colIndex + 1} = vec4(xTexelC${colIndex}.zw, xTexelC${
+                    colIndex + 1}.xy);
+                    `;
               }
-
-              mainLoop += `
-                  xC${colIndex + 1} = vec4(xTexelC${colIndex}.zw, xTexelC${
-                  colIndex + 1}.xy);
-                  `;
             } else {
               // If dilation is 1 and padding is odd, we have already read the
               // texel when constructing the previous x value. Here we can

tfjs-core/src/ops/depthwise_conv2d_test.ts

@@ -717,6 +717,37 @@ describeWithFlags('depthwiseConv2D', ALL_ENVS, () => {
     ]);
   });
 
+  it('input=1x8x8x2,f=3,s=1,d=3,p=valid,chMul=1', async () => {
+    const fSize = 3;
+    const pad = 'valid';
+    const stride = 1;
+    const inDepth = 2;
+    const dilation = 3;
+
+    const x = tf.tensor4d(
+        [
+          1,  1,  2,  2,  3,  3,  4,  4,  5,  5,  6,  6,  7,  7,  8,  8,
+          9,  9,  10, 10, 11, 11, 12, 12, 13, 13, 14, 14, 15, 15, 16, 16,
+          17, 17, 18, 18, 19, 19, 20, 20, 21, 21, 22, 22, 23, 23, 24, 24,
+          25, 25, 26, 26, 27, 27, 28, 28, 29, 29, 30, 30, 31, 31, 32, 32,
+          33, 33, 34, 34, 35, 35, 36, 36, 37, 37, 38, 38, 39, 39, 40, 40,
+          41, 41, 42, 42, 43, 43, 44, 44, 45, 45, 46, 46, 47, 47, 48, 48,
+          49, 49, 50, 50, 51, 51, 52, 52, 53, 53, 54, 54, 55, 55, 56, 56,
+          57, 57, 58, 58, 59, 59, 60, 60, 61, 61, 62, 62, 63, 63, 64, 64
+        ],
+        [1, 8, 8, inDepth]);
+
+    const w = tf.tensor4d(
+        [9, 1, 8, 2, 7, 3, 6, 4, 5, 5, 4, 6, 3, 7, 2, 8, 1, 9],
+        [fSize, fSize, inDepth, 1],
+    );
+    const result = tf.depthwiseConv2d(x, w, stride, pad, 'NHWC', dilation);
+
+    expect(result.shape).toEqual([1, 2, 2, 2]);
+    expectArraysClose(
+        await result.data(), [810, 1710, 855, 1755, 1170, 2070, 1215, 2115]);
+  });
+
   it('Tensor3D is allowed', async () => {
     const fSize = 2;
     const pad = 'same';