If matMul is going to be unpacked, then don't try to avoid expensive reshape in pointwise conv. (tensorflow#1610)

BUG

This PR fixes tensorflow/tfjs#1335

Related: tensorflow/tfjs#1336

Author: annxingyuan

src/kernels/backend_webgl.ts

@@ -1709,9 +1709,19 @@ export class MathBackendWebGL implements KernelBackend {
     // result from 2D to 4D.
     const xShape = x.shape;
     const xTexData = this.texData.get(x.dataId);
-    if (!ENV.get('WEBGL_LAZILY_UNPACK') ||
-        !ENV.get('WEBGL_PACK_BINARY_OPERATIONS') || xShape[2] % 2 === 0 ||
-        !xTexData.isPacked) {
+    const sharedMatMulDim = convInfo.inChannels;
+    const outerShapeX = xShape[0] * xShape[1] * xShape[2];
+    const outerShapeFilter = convInfo.outChannels;
+
+    // TODO: Once reduction ops are packed, batchMatMul will always be packed
+    // and we can remove this condition.
+    const batchMatMulWillBeUnpacked =
+        (outerShapeX === 1 || outerShapeFilter === 1) &&
+        sharedMatMulDim > MATMUL_SHARED_DIM_THRESHOLD;
+    const reshapeWillBeExpensive = xShape[2] % 2 !== 0 && !!xTexData.isPacked;
+
+    if (batchMatMulWillBeUnpacked || !ENV.get('WEBGL_LAZILY_UNPACK') ||
+        !ENV.get('WEBGL_PACK_BINARY_OPERATIONS') || !reshapeWillBeExpensive) {
       const xReshaped =
           this.reshape(
               x, [1, xShape[0] * xShape[1] * xShape[2], convInfo.inChannels]) as

src/ops/conv2d_test.ts

@@ -34,8 +34,8 @@ function generateCaseInputs(totalSizeTensor: number, totalSizeFilter: number) {
   return {input: inp, filter: filt};
 }
 
-describeWithFlags('im2col', PACKED_ENVS, () => {
-  it('should not leak memory', () => {
+describeWithFlags('conv to matmul', PACKED_ENVS, () => {
+  it('im2col should not leak memory', () => {
     const inputDepth = 1;
     const inputShape: [number, number, number] = [2, 2, inputDepth];
     const outputDepth = 1;
@@ -55,6 +55,26 @@ describeWithFlags('im2col', PACKED_ENVS, () => {
 
     expect(endNumBytes - startNumBytes).toEqual(4);
   });
+
+  it('pointwise conv should work when matmul is unpacked', () => {
+    const inputDepth =
+        1001;  // this number must be greater than MATMUL_SHARED_DIM_THRESHOLD
+               // for matmul to be unpacked
+    const inputShape: [number, number, number] = [3, 3, inputDepth];
+    const outputDepth = 1;
+    const fSize = 1;
+    const pad = 'same';
+    const stride: [number, number] = [1, 1];
+
+    let x = tf.randomNormal(inputShape) as tf.Tensor3D;
+    x = x.add(1);  // this packs x so we can test the case where we mistakenly
+                   // want to avoid expensive reshape in pointwise conv2d even
+                   // though matmul is unpacked
+    const w =
+        tf.randomNormal([fSize, fSize, inputDepth, outputDepth]) as tf.Tensor4D;
+
+    expect(() => tf.conv2d(x, w, stride, pad)).not.toThrow();
+  });
 });
 
 describeWithFlags('conv2d', ALL_ENVS, () => {