Implement packed clipping op. (tensorflow#1412)

PERF

- Added `ClipPackedProgram` that implements clipping for packed textures
- Added ENV flag `WEBGL_PACK_CLIP`, defaults to `false`
- Added ENV flag `WEBGL_PACK`, defaults to `false`, that turns all packing-related flags on or off
- Turn off NaN checking if `PROD` flag is on
- Support scalar sampling in packed ops
- Partially support packed `get${Var}AtOutCoords` and broadcasting - broadcasting only occurs when the innermost dimensions are the same and the vector to be broadcast is rank 0 or 1.

Author: annxingyuan

src/environment.ts

@@ -308,14 +308,18 @@ export class Environment {
       return isChrome();
     } else if (feature === 'WEBGL_CPU_FORWARD') {
       return true;
-    } else if (feature === 'WEBGL_PACK_BATCHNORMALIZATION') {
+    } else if (feature === 'WEBGL_PACK') {
       return false;
+    } else if (feature === 'WEBGL_PACK_BATCHNORMALIZATION') {
+      return this.get('WEBGL_PACK');
+    } else if (feature === 'WEBGL_PACK_CLIP') {
+      return this.get('WEBGL_PACK');
     } else if (feature === 'WEBGL_PACK_DEPTHWISECONV') {
-      return false;
+      return this.get('WEBGL_PACK');
     } else if (feature === 'WEBGL_LAZILY_UNPACK') {
-      return false;
+      return this.get('WEBGL_PACK');
     } else if (feature === 'WEBGL_CONV_IM2COL') {
-      return false;
+      return this.get('WEBGL_PACK');
     } else if (feature === 'WEBGL_PAGING_ENABLED') {
       return this.get('IS_BROWSER') && !this.get('PROD');
     } else if (feature === 'WEBGL_MAX_TEXTURE_SIZE') {

src/environment_util.ts

@@ -28,8 +28,12 @@ export interface Features {
   'WEBGL_LAZILY_UNPACK'?: boolean;
   // Whether the WebGL backend will sometimes forward ops to the CPU.
   'WEBGL_CPU_FORWARD'?: boolean;
+  // Whether to turn all packing related flags on.
+  'WEBGL_PACK'?: boolean;
   // Whether we will pack the batchnormalization op.
   'WEBGL_PACK_BATCHNORMALIZATION'?: boolean;
+  // Whether we will pack the clipping op.
+  'WEBGL_PACK_CLIP'?: boolean;
   // Whether we pack the depthwise convolution op.
   'WEBGL_PACK_DEPTHWISECONV'?: boolean;
   // Whether we will use the im2col algorithm to speed up convolutions.
@@ -91,7 +95,9 @@ export const URL_PROPERTIES: URLProperty[] = [
   {name: 'IS_BROWSER', type: Type.BOOLEAN},
   {name: 'WEBGL_LAZILY_UNPACK', type: Type.BOOLEAN},
   {name: 'WEBGL_CPU_FORWARD', type: Type.BOOLEAN},
+  {name: 'WEBGL_PACK', type: Type.BOOLEAN},
   {name: 'WEBGL_PACK_BATCHNORMALIZATION', type: Type.BOOLEAN},
+  {name: 'WEBGL_PACK_CLIP', type: Type.BOOLEAN},
   {name: 'WEBGL_PACK_DEPTHWISECONV', type: Type.BOOLEAN},
   {name: 'WEBGL_CONV_IM2COL', type: Type.BOOLEAN},
   {name: 'WEBGL_MAX_TEXTURE_SIZE', type: Type.NUMBER},

src/kernels/backend_webgl.ts

@@ -51,6 +51,7 @@ import {BinaryOpComplexProgram} from './webgl/binaryop_complex_gpu';
 import * as binaryop_gpu from './webgl/binaryop_gpu';
 import {BinaryOpProgram} from './webgl/binaryop_gpu';
 import {ClipProgram} from './webgl/clip_gpu';
+import {ClipPackedProgram} from './webgl/clip_packed_gpu';
 import {ComplexAbsProgram} from './webgl/complex_abs_gpu';
 import {ConcatProgram} from './webgl/concat_gpu';
 import {Conv2DDerFilterProgram, Conv2DDerInputProgram} from './webgl/conv_backprop_gpu';
@@ -366,9 +367,11 @@ export class MathBackendWebGL implements KernelBackend {
     const {shape, dtype, texture, texShape} = this.texData.get(dataId);
     if (ENV.get('WEBGL_DOWNLOAD_FLOAT_ENABLED')) {
       if (this.texData.get(dataId).isPacked) {
-        const batch = util.sizeFromShape(shape.slice(0, shape.length - 2));
-        const rows = shape.length > 1 ? shape[shape.length - 2] : 1;
-        const cols = shape[shape.length - 1];
+        const batch = this.getBatchDim(shape);
+        let rows = 1, cols = 1;
+        if (shape.length) {
+          [rows, cols] = this.getRowsCols(shape);
+        }
         return this.gpgpu.downloadMatrixFromPackedTexture(
             texture, batch, rows, cols, texShape[0], texShape[1]);
       } else {
@@ -1337,7 +1340,12 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   clip<T extends Tensor>(x: T, min: number, max: number): T {
-    const program = new ClipProgram(x.shape, min, max);
+    let program;
+    if (ENV.get('WEBGL_PACK_CLIP')) {
+      program = new ClipPackedProgram(x.shape, min, max);
+    } else {
+      program = new ClipProgram(x.shape, min, max);
+    }
     return this.compileAndRun(program, [x]) as T;
   }
 

src/kernels/webgl/batchnorm_packed_gpu.ts

@@ -33,34 +33,29 @@ export class BatchNormPackedProgram implements GPGPUProgram {
     broadcast_util.assertAndGetBroadcastShape(xShape, meanShape);
     broadcast_util.assertAndGetBroadcastShape(xShape, varianceShape);
 
-    const meanSnippet = broadcastSample('mean', meanShape.length);
-    const varianceSnippet = broadcastSample('variance', varianceShape.length);
-
-    let offsetSnippet = 'vec4 offset = vec4(0.0)';
+    let offsetSnippet = 'vec4(0.0)';
     if (offsetShape != null) {
       broadcast_util.assertAndGetBroadcastShape(xShape, offsetShape);
       this.variableNames.push('offset');
-      offsetSnippet = broadcastSample('offset', offsetShape.length);
+      offsetSnippet = 'getOffsetAtOutCoords()';
     }
 
-    let scaleSnippet = 'vec4 scale = vec4(1.0)';
+    let scaleSnippet = 'vec4(1.0)';
     if (scaleShape != null) {
       broadcast_util.assertAndGetBroadcastShape(xShape, scaleShape);
       this.variableNames.push('scale');
-      scaleSnippet = broadcastSample('scale', scaleShape.length);
+      scaleSnippet = 'getScaleAtOutCoords()';
     }
 
     this.outputShape = xShape;
     this.userCode = `
       void main() {
-        ivec4 rc = getOutputCoords();
-
-        ${offsetSnippet};
-        ${scaleSnippet};
+        vec4 offset = ${offsetSnippet};
+        vec4 scale = ${scaleSnippet};
 
-        vec4 x = getX(rc.x, rc.y, rc.z, rc.w);
-        ${meanSnippet};
-        ${varianceSnippet};
+        vec4 x = getXAtOutCoords();
+        vec4 mean = getMeanAtOutCoords();
+        vec4 variance = getVarianceAtOutCoords();
 
         vec4 inv = scale * inversesqrt(variance + vec4(${varianceEpsilon}));
 
@@ -69,14 +64,3 @@ export class BatchNormPackedProgram implements GPGPUProgram {
     `;
   }
 }
-
-function broadcastSample(texName: string, rank: number): string {
-  const texSampler = `get${texName.charAt(0).toUpperCase()}${texName.slice(1)}`;
-  if (rank === 1) {
-    return `
-      vec4 ${texName}Sample = ${texSampler}(rc.w);
-      vec4 ${texName} = vec4(${texName}Sample.xy, ${texName}Sample.xy);
-    `;
-  }
-  return `vec4 ${texName} = ${texSampler}(rc.x, rc.y, rc.z, rc.w)`;
-}

src/kernels/webgl/clip_packed_gpu.ts

@@ -0,0 +1,41 @@
+/**
+ * @license
+ * Copyright 2018 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GPGPUProgram} from './gpgpu_math';
+
+export class ClipPackedProgram implements GPGPUProgram {
+  variableNames = ['A'];
+  usesPackedTextures = true;
+  userCode: string;
+  outputShape: number[];
+
+  constructor(aShape: number[], min: number, max: number) {
+    this.outputShape = aShape;
+    this.userCode = `
+      void main() {
+        vec4 value = getAAtOutCoords();
+
+        if (hasNaN(value)) {
+          setOutput(value);
+          return;
+        }
+
+        setOutput(clamp(value, vec4(${min}), vec4(${max})));
+      }
+    `;
+  }
+}

src/kernels/webgl/pack_gpu.ts

@@ -33,28 +33,36 @@ export class PackProgram implements GPGPUProgram {
     this.outputShape = outputShape;
     const rank = outputShape.length;
 
-    const channels = getChannels('rc', rank);
-    const dtype = getCoordsDataType(rank);
-    const outOfBoundsCondition =
-        getOutOfBoundsCondition(rank, outputShape, channels);
-    const setup = getSetup(
-        rank, outputShape[outputShape.length - 1],
-        outputShape[outputShape.length - 2], channels);
-    const output = getOutput(outputShape, channels);
-
-    this.userCode = `
-      void main() {
-        ${dtype} rc = getOutputCoords();
-
-        if(${outOfBoundsCondition}) {
-          gl_FragColor = vec4(0);
-        } else {
-          ${setup}
-
-          setOutput(vec4(${output}));
+    if (rank === 0) {
+      this.userCode = `
+        void main() {
+          setOutput(vec4(getA(), 0., 0., 0.));
         }
-      }
-    `;
+      `;
+    } else {
+      const channels = getChannels('rc', rank);
+      const dtype = getCoordsDataType(rank);
+      const outOfBoundsCondition =
+          getOutOfBoundsCondition(rank, outputShape, channels);
+      const setup = getSetup(
+          rank, outputShape[outputShape.length - 1],
+          outputShape[outputShape.length - 2], channels);
+      const output = getOutput(outputShape, channels);
+
+      this.userCode = `
+        void main() {
+          ${dtype} rc = getOutputCoords();
+
+          if(${outOfBoundsCondition}) {
+            setOutput(vec4(0));
+          } else {
+            ${setup}
+
+            setOutput(vec4(${output}));
+          }
+        }
+      `;
+    }
   }
 }