Pack unary ops.

src/kernels/backend_webgl.ts

@@ -106,6 +106,7 @@ import {TransposeProgram} from './webgl/transpose_gpu';
 import * as unary_op from './webgl/unaryop_gpu';
 import {UnaryOpProgram} from './webgl/unaryop_gpu';
 import * as unary_packed_op from './webgl/unaryop_packed_gpu';
+import {UnaryOpPackedProgram} from './webgl/unaryop_packed_gpu';
 import {UnpackProgram} from './webgl/unpack_gpu';
 import * as webgl_util from './webgl/webgl_util';
 import {whereImpl} from './where_impl';
@@ -1440,7 +1441,12 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   exp<T extends Tensor>(x: T): T {
-    const program = new UnaryOpProgram(x.shape, unary_op.EXP);
+    let program: UnaryOpProgram | UnaryOpPackedProgram;
+    if(ENV.get('WEBGL_PACK')) {
+      program = new UnaryOpPackedProgram(x.shape, unary_op.EXP);
+    } else {
+      program = new UnaryOpProgram(x.shape, unary_op.EXP);
+    }
     return this.compileAndRun(program, [x]) as T;
   }
 
@@ -1450,7 +1456,12 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   log<T extends Tensor>(x: T): T {
-    const program = new UnaryOpProgram(x.shape, unary_op.LOG);
+    let program: UnaryOpProgram | UnaryOpPackedProgram;
+    if(ENV.get('WEBGL_PACK')) {
+      program = new UnaryOpPackedProgram(x.shape, unary_packed_op.LOG);
+    } else {
+      program = new UnaryOpProgram(x.shape, unary_op.LOG);
+    }
     const customSetup = program.getCustomSetupFunc();
     return this.compileAndRun(program, [x], null, customSetup) as T;
   }
@@ -1481,7 +1492,12 @@ export class MathBackendWebGL implements KernelBackend {
   }
 
   relu<T extends Tensor>(x: T): T {
-    const program = new UnaryOpProgram(x.shape, unary_op.RELU);
+    let program: UnaryOpProgram | UnaryOpPackedProgram;
+    if(ENV.get('WEBGL_PACK')) {
+      program = new UnaryOpPackedProgram(x.shape, unary_packed_op.RELU);
+    } else {
+      program = new UnaryOpProgram(x.shape, unary_op.RELU);
+    }
     return this.compileAndRun(program, [x]) as T;
   }
 

src/kernels/webgl/shader_compiler.ts

@@ -504,15 +504,15 @@ function getOutputPackedNDCoords(
       int b${b} = index / ${texelsInBatchN};
       index -= b${b} * ${texelsInBatchN};
     ` + batches;
-    coords = `b${b}, ` + coords; 
+    coords = `b${b}, ` + coords;
   }
 
   return `
     ivec${shape.length} getOutputCoords() {
       ivec2 resTexRC = ivec2(resultUV.yx *
                              vec2(${packedTexShape[0]}, ${packedTexShape[1]}));
       int index = resTexRC.x * ${packedTexShape[1]} + resTexRC.y;
-      
+
       ${batches}
 
       int b = index / ${texelsInBatch};
@@ -984,7 +984,7 @@ function getSampler3D(inputInfo: InputInfo): string {
   `;
 }
 
-function getPackedSamplerND(inputInfo: InputInfo): string {  
+function getPackedSamplerND(inputInfo: InputInfo): string {
   const shape = inputInfo.shapeInfo.logicalShape;
   const rank = shape.length;
   const texName = inputInfo.name;
@@ -1001,7 +1001,7 @@ function getPackedSamplerND(inputInfo: InputInfo): string {
   let index = `b * ${texelsInBatch} + (row / 2) * ${valuesPerRow} + (col / 2)`;
   for (let b = 2; b < rank - 1; b++) {
     params = `int b${b}, ` + params;
-    texelsInBatch *= shape[rank - b - 1]; 
+    texelsInBatch *= shape[rank - b - 1];
     index = `b${b} * ${texelsInBatch} + ` + index;
   }
   const glsl = getGlslDifferences();
@@ -1011,7 +1011,7 @@ function getPackedSamplerND(inputInfo: InputInfo): string {
       int texR = index / ${texNumC};
       int texC = index - texR * ${texNumC};
       vec2 uv = (vec2(texC, texR) + halfCR) / vec2(${texNumC}, ${texNumR});
-      return ${glsl.texture2D}(${texName}, uv);      
+      return ${glsl.texture2D}(${texName}, uv);
     }
   `;
 }

src/kernels/webgl/unaryop_packed_gpu.ts

@@ -15,9 +15,70 @@
  * =============================================================================
  */
 
-const CHECK_NAN_SNIPPET = `if (hasNaN(x)) return x;`;
+import {GPGPUContext} from './gpgpu_context';
+import {GPGPUProgram} from './gpgpu_math';
 
 export const LINEAR = `return x;`;
 
-export const RELU =
-    CHECK_NAN_SNIPPET + `return x * vec4(greaterThanEqual(x, vec4(0.0)));`;
+export const LOG = `
+  vec4 result = log(x);
+  vec4 isNaN = vec4(lessThan(x, vec4(0.0)));
+  result.r = isNaN.r == 1.0 ? NAN : result.r;
+  result.g = isNaN.g == 1.0 ? NAN : result.g;
+  result.b = isNaN.b == 1.0 ? NAN : result.b;
+  result.a = isNaN.a == 1.0 ? NAN : result.a;
+
+  return result;
+`;
+
+export const RELU = `
+  vec4 result = x * vec4(greaterThanEqual(x, vec4(0.0)));
+
+  result.r = isNaN(x.r) ? x.r : result.r;
+  result.g = isNaN(x.g) ? x.g : result.g;
+  result.b = isNaN(x.b) ? x.b : result.b;
+  result.a = isNaN(x.a) ? x.a : result.a;
+
+  return result;
+`;
+
+export class UnaryOpPackedProgram implements GPGPUProgram {
+  variableNames = ['A'];
+  userCode: string;
+  outputShape: number[];
+  usesPackedTextures = true;
+
+  // Caching uniform location for speed.
+  startLoc: WebGLUniformLocation;
+
+  constructor(aShape: number[], opSnippet: string) {
+    this.outputShape = aShape;
+    this.userCode = `
+      uniform float NAN;
+      vec4 unaryOperation(vec4 x) {
+        ${opSnippet}
+      }
+
+      void main() {
+        vec4 x = getAAtOutCoords();
+        vec4 y = unaryOperation(x);
+
+        setOutput(y);
+      }
+    `;
+  }
+
+  getCustomSetupFunc() {
+    return (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => {
+      if (this.startLoc == null) {
+        this.startLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'NAN');
+        if (this.startLoc == null) {
+          // This means the compiler has optimized and realized it doesn't need
+          // the uniform.
+          return;
+        }
+      }
+      gpgpu.gl.uniform1f(this.startLoc, NaN);
+    };
+  }
+}