cache attribute locations in gpgpu binary (#169)

* cache attribute locations in gpgpu binary

* address review

* fix typo

Author: dsmilkov

src/math/webgl/gpgpu_context.ts

@@ -198,6 +198,13 @@ export class GPGPUContext {
         this.gl, program, uniformName);
   }
 
+  public getAttributeLocation(program: WebGLProgram, attribute: string):
+      number {
+    this.throwIfDisposed();
+    return webgl_util.callAndCheck(
+        this.gl, () => this.gl.getAttribLocation(program, attribute));
+  }
+
   public getUniformLocationNoThrow(program: WebGLProgram, uniformName: string):
       WebGLUniformLocation {
     this.throwIfDisposed();
@@ -247,12 +254,12 @@ export class GPGPUContext {
     webgl_util.validateFramebuffer(this.gl);
   }
 
-  public executeProgram() {
+  public executeProgram(attribLocations?: {[name: string]: number}) {
     this.throwIfDisposed();
     this.throwIfNoProgram();
     const gl = this.gl;
     gpgpu_util.bindVertexProgramAttributeStreams(
-        gl, this.program, this.vertexBuffer);
+        gl, this.program, this.vertexBuffer, attribLocations);
     if (this.autoDebugValidate) {
       this.debugValidate();
     }

src/math/webgl/gpgpu_math.ts

@@ -22,6 +22,8 @@ import {GPGPUContext} from './gpgpu_context';
 import * as shader_compiler from './shader_compiler';
 import {ShapeInfo} from './shader_compiler';
 
+const ATTRIBUTE_NAMES = ['uv', 'clipSpacePos'];
+
 export interface GPGPUProgram {
   variableNames: string[];
   outputShape: number[];
@@ -34,6 +36,7 @@ export interface GPGPUBinary {
   webGLProgram: WebGLProgram;
   program: GPGPUProgram;
   uniformLocations: {[name: string]: WebGLUniformLocation};
+  attributeLocations: {[name: string]: number};
   gpgpu: GPGPUContext;
   source: string;
   inShapeInfos: ShapeInfo[];
@@ -68,12 +71,18 @@ export function compileProgram<T extends NDArray, K extends NDArray>(
     uniformLocations[uniformName] =
         gpgpu.getUniformLocation(webGLProgram, uniformName);
   }
+  const attributeLocations: {[name: string]: number} = {};
+  ATTRIBUTE_NAMES.forEach(attribute => {
+    attributeLocations[attribute] =
+        gpgpu.getAttributeLocation(webGLProgram, attribute);
+  });
 
   return {
     program,
     source,
     webGLProgram,
     uniformLocations,
+    attributeLocations,
     gpgpu,
     inShapeInfos,
     outShapeInfo
@@ -127,7 +136,7 @@ export function runProgram<T extends NDArray, K extends NDArray>(
   if (customSetup != null) {
     customSetup(gpgpu, binary.webGLProgram);
   }
-  gpgpu.executeProgram();
+  gpgpu.executeProgram(binary.attributeLocations);
 }
 
 export function makeShaderKey(

src/math/webgl/gpgpu_util.ts

@@ -156,26 +156,18 @@ export function createPackedMatrixTexture(
 }
 
 export function bindVertexProgramAttributeStreams(
-    gl: WebGLRenderingContext, program: WebGLProgram,
-    vertexBuffer: WebGLBuffer) {
+    gl: WebGLRenderingContext, program: WebGLProgram, vertexBuffer: WebGLBuffer,
+    attribLocations?: {[name: string]: number}) {
   const posOffset = 0;               // x is the first buffer element
   const uvOffset = 3 * 4;            // uv comes after [x y z]
   const stride = (3 * 4) + (2 * 4);  // xyz + uv, each entry is 4-byte float.
   webgl_util.callAndCheck(
       gl, () => gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer));
   webgl_util.bindVertexBufferToProgramAttribute(
-      gl, program, 'clipSpacePos', vertexBuffer, 3, stride, posOffset);
-  try {
-    webgl_util.bindVertexBufferToProgramAttribute(
-        gl, program, 'uv', vertexBuffer, 2, stride, uvOffset);
-  } catch (e) {
-    // Programs with 1x1 output textures don't use the uv attribute.
-    // This can cause the shader linker to dead-strip it, so we shouldn't
-    // complain or fail if it's not present.
-    if (!e.hasOwnProperty('namedVertexAttributeNotFound')) {
-      throw e;
-    }
-  }
+      gl, program, 'clipSpacePos', vertexBuffer, 3, stride, posOffset,
+      attribLocations);
+  webgl_util.bindVertexBufferToProgramAttribute(
+      gl, program, 'uv', vertexBuffer, 2, stride, uvOffset, attribLocations);
 }
 
 export function uploadPixelDataToTexture(

src/math/webgl/webgl_util.ts

@@ -257,14 +257,17 @@ export function createFramebuffer(gl: WebGLRenderingContext): WebGLFramebuffer {
 export function bindVertexBufferToProgramAttribute(
     gl: WebGLRenderingContext, program: WebGLProgram, attribute: string,
     buffer: WebGLBuffer, arrayEntriesPerItem: number, itemStrideInBytes: number,
-    itemOffsetInBytes: number) {
-  const loc = gl.getAttribLocation(program, attribute);
+    itemOffsetInBytes: number, attribLocations?: {[name: string]: number}) {
+  let loc = -1;
+  if ((attribLocations != null) && (attribute in attribLocations)) {
+    loc = attribLocations[attribute];
+  } else {
+    loc = gl.getAttribLocation(program, attribute);
+  }
   if (loc === -1) {
-    const error = new Error(
-        'Unable to get attribute "' + attribute + '" on WebGLProgram.');
-    // tslint:disable-next-line:no-any
-    (error as any).namedVertexAttributeNotFound = attribute;
-    throw error;
+    // The GPU compiler decided to strip out this attribute because it's unused,
+    // thus no need to bind.
+    return;
   }
   callAndCheck(gl, () => gl.bindBuffer(gl.ARRAY_BUFFER, buffer));
   callAndCheck(