[webgl] Reduce JS map lookups in webgl backend

tfjs-backend-webgl/src/backend_webgl.ts

@@ -1276,24 +1276,20 @@ export class MathBackendWebGL extends KernelBackend {
   }
 
   getUniformLocations() {
-    for (const [, binary] of Object.entries(this.binaryCache)) {
+    for (const binary of Object.values(this.binaryCache)) {
       const {
-        uniformLocations,
+        variablesLocations,
         customUniformLocations,
         infLoc,
         nanLoc,
-        inShapesLocations,
-        inTexShapesLocations,
         outShapeLocation,
         outShapeStridesLocation,
         outTexShapeLocation
       } = getUniformLocations(this.gpgpu, binary.program, binary.webGLProgram);
-      binary.uniformLocations = uniformLocations;
+      binary.variablesLocations = variablesLocations;
       binary.customUniformLocations = customUniformLocations;
       binary.infLoc = infLoc;
       binary.nanLoc = nanLoc;
-      binary.inShapesLocations = inShapesLocations;
-      binary.inTexShapesLocations = inTexShapesLocations;
       binary.outShapeLocation = outShapeLocation;
       binary.outShapeStridesLocation = outShapeStridesLocation;
       binary.outTexShapeLocation = outTexShapeLocation;

tfjs-backend-webgl/src/gpgpu_math.ts

@@ -46,34 +46,31 @@ export interface GPGPUProgram {
       Array<{name: string; arrayIndex?: number; type: UniformType;}>;
 }
 
-export interface GPGPUBinary {
+export interface GPGPUBinary extends GPGPUBinaryLocations {
   webGLProgram: GPGPUContextProgram;
   program: GPGPUProgram;
-  uniformLocations: {[name: string]: WebGLUniformLocation};
-  customUniformLocations?: WebGLUniformLocation[];
   source: string;
   fragmentShader: WebGLShader;
   inShapeInfos: ShapeInfo[];
   outShapeInfo: ShapeInfo;
-  infLoc: WebGLUniformLocation;
-  nanLoc: WebGLUniformLocation;
-  inShapesLocations?: {[name: string]: WebGLUniformLocation};
-  inTexShapesLocations?: {[name: string]: WebGLUniformLocation};
-  outShapeLocation?: WebGLUniformLocation;
-  outShapeStridesLocation?: WebGLUniformLocation;
-  outTexShapeLocation?: WebGLUniformLocation;
 }
 
 export interface GPGPUBinaryLocations {
-  uniformLocations: {[name: string]: WebGLUniformLocation};
   customUniformLocations?: WebGLUniformLocation[];
   infLoc: WebGLUniformLocation;
   nanLoc: WebGLUniformLocation;
-  inShapesLocations?: {[name: string]: WebGLUniformLocation};
-  inTexShapesLocations?: {[name: string]: WebGLUniformLocation};
   outShapeLocation?: WebGLUniformLocation;
   outShapeStridesLocation?: WebGLUniformLocation;
   outTexShapeLocation?: WebGLUniformLocation;
+  variablesLocations?: GPGPUVariableLocations[];
+}
+
+export interface GPGPUVariableLocations {
+  name: string;
+  uniform: WebGLUniformLocation;
+  offset: WebGLUniformLocation;
+  shape?: WebGLUniformLocation;
+  texShape?: WebGLUniformLocation;
 }
 
 export interface TensorData {
@@ -131,12 +128,10 @@ export function compileProgram<T extends Tensor, K extends Tensor>(
       webGLProgram,
       inShapeInfos,
       outShapeInfo,
-      uniformLocations: null,
+      variablesLocations: null,
       customUniformLocations: null,
       infLoc: null,
       nanLoc: null,
-      inShapesLocations: null,
-      inTexShapesLocations: null,
       outShapeLocation: null,
       outShapeStridesLocation: null,
       outTexShapeLocation: null
@@ -147,9 +142,7 @@ export function compileProgram<T extends Tensor, K extends Tensor>(
 export function getUniformLocations(
     gpgpu: GPGPUContext, program: GPGPUProgram,
     webGLProgram: WebGLProgram): GPGPUBinaryLocations {
-  const uniformLocations: {[name: string]: WebGLUniformLocation} = {};
-  const inShapesLocations: {[name: string]: WebGLUniformLocation} = {};
-  const inTexShapesLocations: {[name: string]: WebGLUniformLocation} = {};
+  const variablesLocations: GPGPUVariableLocations[] = [];
   const customUniformLocations: WebGLUniformLocation[] = [];
   let outShapeLocation: WebGLUniformLocation;
   let outTexShapeLocation: WebGLUniformLocation;
@@ -165,18 +158,21 @@ export function getUniformLocations(
 
   // Add user-defined uniforms
   const shouldThrow = false;
-  for (let i = 0; i < program.variableNames.length; i++) {
-    const varName = program.variableNames[i];
-    uniformLocations[varName] =
-        gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow);
-    uniformLocations[`offset${varName}`] =
-        gpgpu.getUniformLocation(webGLProgram, `offset${varName}`, shouldThrow);
+  for (const varName of program.variableNames) {
+    const varLocs: GPGPUVariableLocations = {
+      name: varName,
+      uniform: gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow),
+      offset: gpgpu.getUniformLocation(
+          webGLProgram, `offset${varName}`, shouldThrow),
+    };
     if (program.enableShapeUniforms) {
-      inShapesLocations[`${varName}Shape`] = gpgpu.getUniformLocation(
+      varLocs.shape = gpgpu.getUniformLocation(
           webGLProgram, `${varName}Shape`, shouldThrow);
-      inTexShapesLocations[`${varName}TexShape`] = gpgpu.getUniformLocation(
+      varLocs.texShape = gpgpu.getUniformLocation(
           webGLProgram, `${varName}TexShape`, shouldThrow);
     }
+
+    variablesLocations.push(varLocs);
   }
 
   if (program.enableShapeUniforms) {
@@ -189,19 +185,17 @@ export function getUniformLocations(
   }
 
   if (program.customUniforms) {
-    program.customUniforms.forEach((d, i) => {
-      customUniformLocations[i] =
-          gpgpu.getUniformLocation(webGLProgram, d.name, shouldThrow);
-    });
+    for (const d of program.customUniforms) {
+      customUniformLocations.push(
+          gpgpu.getUniformLocation(webGLProgram, d.name, shouldThrow));
+    }
   }
 
   return {
-    uniformLocations,
+    variablesLocations,
     customUniformLocations,
     infLoc,
     nanLoc,
-    inShapesLocations,
-    inTexShapesLocations,
     outShapeLocation,
     outShapeStridesLocation,
     outTexShapeLocation
@@ -271,12 +265,14 @@ export function runProgram<T extends Tensor, K extends Tensor>(
   }
 
   // Set user-defined inputs
-  inputs.forEach((input, i) => {
-    const varName = binary.program.variableNames[i];
-    const varLoc = binary.uniformLocations[varName];
-    const varOffsetLoc = binary.uniformLocations[`offset${varName}`];
-    const varShapeLoc = binary.inShapesLocations[`${varName}Shape`];
-    const varTexShapeLoc = binary.inTexShapesLocations[`${varName}TexShape`];
+  for (let i = 0; i < inputs.length; ++i) {
+    const input = inputs[i];
+    const {
+      uniform: varLoc,
+      offset: varOffsetLoc,
+      shape: varShapeLoc,
+      texShape: varTexShapeLoc,
+    } = binary.variablesLocations[i];
 
     if (varShapeLoc) {
       const {uniformShape} = shader_compiler.getUniformInfoFromShape(
@@ -298,14 +294,15 @@ export function runProgram<T extends Tensor, K extends Tensor>(
           break;
       }
     }
+
     if (varTexShapeLoc) {
       gpgpu.gl.uniform2i(
           varTexShapeLoc, input.texData.texShape[0], input.texData.texShape[1]);
     }
 
     if (varLoc == null) {
       // The compiler inferred that this variable is not used in this shader.
-      return;
+      continue;
     }
 
     if (input.isUniform) {
@@ -319,7 +316,7 @@ export function runProgram<T extends Tensor, K extends Tensor>(
         }
         gpgpu.gl.uniform1fv(varLoc, vals);
       }
-      return;
+      continue;
     }
 
     // If the input was sliced, upload the flat offset index.
@@ -328,7 +325,7 @@ export function runProgram<T extends Tensor, K extends Tensor>(
     }
 
     gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i);
-  });
+  }
 
   const outShapeLoc = binary.outShapeLocation;
   if (outShapeLoc) {
@@ -375,7 +372,8 @@ export function runProgram<T extends Tensor, K extends Tensor>(
   }
 
   if (binary.program.customUniforms && customUniformValues) {
-    binary.program.customUniforms.forEach((d, i) => {
+    for (let i = 0; i < binary.program.customUniforms.length; ++i) {
+      const d = binary.program.customUniforms[i];
       const customLoc = binary.customUniformLocations[i];
       const customValue = customUniformValues[i];
       if (d.type === 'float') {
@@ -397,7 +395,7 @@ export function runProgram<T extends Tensor, K extends Tensor>(
       } else {
         throw Error(`uniform type ${d.type} is not supported yet.`);
       }
-    });
+    }
   }
   gpgpu.executeProgram();
 }

tfjs-backend-webgl/src/texture_manager.ts

@@ -25,13 +25,13 @@ export class TextureManager {
   private numUsedTextures = 0;
   private numFreeTextures = 0;
   private _numBytesAllocated = 0;
-  private _numBytesFree = 0;  // How many bytes that have been allocated
-                              // are available for reuse.
-  private freeTextures: {[shape: string]: Texture[]} = {};
+  // Number of bytes that have been allocated and available for reuse.
+  private _numBytesFree = 0;
+  private freeTextures: Record<string, Texture[]> = {};
+  private usedTextures: Record<string, Texture[]> = {};
   private logEnabled = false;
-  private usedTextures: {[shape: string]: Texture[]} = {};
 
-  constructor(private gpgpu: GPGPUContext) {}
+  constructor(private readonly gpgpu: GPGPUContext) {}
 
   acquireTexture(
       shapeRC: [number, number], usage: TextureUsage,
@@ -55,7 +55,7 @@ export class TextureManager {
       this.numUsedTextures++;
       this._numBytesFree -= texBytes;
       this.log();
-      const newTexture = this.freeTextures[shapeKey].shift();
+      const newTexture = this.freeTextures[shapeKey].pop();
       this.usedTextures[shapeKey].push(newTexture);
       return newTexture;
     }
@@ -117,13 +117,14 @@ export class TextureManager {
     this.numUsedTextures--;
 
     const texList = this.usedTextures[shapeKey];
-    const texIndex = texList.indexOf(texture);
-    if (texIndex < 0) {
+    const texIndex = texList && texList.indexOf(texture);
+    if (texIndex == null || texIndex < 0) {
       throw new Error(
           'Cannot release a texture that was never provided by this ' +
           'texture manager');
     }
-    texList.splice(texIndex, 1);
+    texList[texIndex] = texList[texList.length - 1];
+    texList.pop();
     this.log();
   }
 
@@ -172,6 +173,7 @@ export class TextureManager {
         this.gpgpu.deleteMatrixTexture(tex.texture);
       });
     }
+    // TODO: Assign non-null value (empty object) to textures after disposed.
     this.freeTextures = null;
     this.usedTextures = null;
     this.numUsedTextures = 0;