webgl: various fixes for stability and perf (#34)

* renaming TextureManager

* adding --printMatches to verify accuracy visually

* Ensuring textures are released at the end of each inference

* binding Attributes(vertices) only once

* fixing a bug with Reshape operator and caching texturedata

* removing redundant call

* changing number[] to ReadonlyArray<number>

benchmark/package.json

@@ -6,9 +6,9 @@
   "scripts": {
     "build": "webpack --config ./webpack.conf.js --mode production",
     "build-debug": "webpack --config ./webpack.conf.js --mode development",
-    "test": "karma start --browsers ChromeTest --single-run",
+    "test": "karma start --browsers ChromeTest --single-run  --printMatches",
     "test-debug": "karma start --browsers ChromeDebug  --printMatches",
-    "test-edge": "karma start --browsers Edge --single-run"
+    "test-edge": "karma start --browsers Edge --single-run  --printMatches"
   },
   "keywords": [
     "resnet",

lib/backends/webgl/base-webgl-context.ts

@@ -229,6 +229,13 @@ export abstract class BaseWebGLContext implements WebGLContext, Disposable {
         throw new Error(`Invalid dataType: ${dataType}`);
     }
   }
+  clearActiveTextures(): void {
+    const gl = this.gl;
+    for (let unit = 0; unit < this.maxTextureImageUnits; ++unit) {
+      gl.activeTexture(gl.TEXTURE0 + unit);
+      gl.bindTexture(gl.TEXTURE_2D, null);
+    }
+  }
   protected queryVitalParameters(): void {
     const gl = this.gl;
     this.maxCombinedTextureImageUnits = gl.getParameter(gl.MAX_COMBINED_TEXTURE_IMAGE_UNITS);

lib/backends/webgl/inference-handler.ts

@@ -10,8 +10,8 @@ import {WebGLBackend} from '../backend-webgl';
 import {ProgramManager} from './program-manager';
 import {WebGLSessionHandler} from './session-handler';
 import {TextureData, TextureLayout} from './texture-data';
+import {TextureHelper} from './texture-helper';
 import {WidthHeightPrefs} from './texture-layout-strategy';
-import {TextureManager} from './texture-manager';
 import {getPackedShape} from './utils';
 
 /**
@@ -20,18 +20,20 @@ import {getPackedShape} from './utils';
  * Throughout WebGL backend operations TextureData is used as the data carrier
  */
 export class WebGLInferenceHandler implements InferenceHandler {
-  textureManager: TextureManager;
+  textureHelper: TextureHelper;
   programManager: ProgramManager;
   private tensorToTexture: Map<Tensor, TextureData>;
   private textureToTensor: Map<TextureData, Tensor>;
   constructor(public backend: WebGLBackend, public session: WebGLSessionHandler) {
-    this.textureManager = session.textureManager;
+    this.textureHelper = session.textureHelper;
     this.programManager = session.programManager;
     this.tensorToTexture = new Map();
     this.textureToTensor = new Map();
   }
   protected lookupTextureData(tensor: Tensor): TextureData|undefined {
-    return this.session.isInitializer(tensor) ? this.session.getTextureData(tensor) : this.tensorToTexture.get(tensor);
+    const isInitializer = this.session.isInitializer(tensor);
+    Logger.verbose('InferenceHandler', `tensor was an initializer; returning TextureData from session cache`);
+    return isInitializer ? this.session.getTextureData(tensor) : this.tensorToTexture.get(tensor);
   }
   getOrCreate(tensor: Tensor, layout?: TextureLayout): TextureData {
     let td = this.lookupTextureData(tensor);
@@ -64,7 +66,7 @@ export class WebGLInferenceHandler implements InferenceHandler {
     if (!tensor) {
       Logger.verbose('InferenceHandler', `Creating new Tensor from texture data: [${td.unpackedShape}]`);
       tensor = new Tensor(td.unpackedShape, td.dataType, (id: Tensor.Id) => {
-        const values = this.textureManager.readTexture(td, td.dataType, td.channels);
+        const values = this.textureHelper.readTexture(td, td.dataType, td.channels);
         return values;
       });
       this.setTextureData(tensor, td);
@@ -82,20 +84,21 @@ export class WebGLInferenceHandler implements InferenceHandler {
         channels === 1 ? tensor.dims.slice() : getPackedShape(tensor.dims.slice()), channels, unpackedShape);
   }
   dispose(): void {
-    this.tensorToTexture.forEach(td => this.textureManager.saveTexture(td.texture, [td.width, td.height]));
+    this.textureHelper.clearActiveTextures();
+    this.tensorToTexture.forEach(td => this.textureHelper.releaseTexture(td.texture));
     this.tensorToTexture = new Map();
     this.textureToTensor = new Map();
   }
   createTextureData(
       dataType: Tensor.DataType, shape: ReadonlyArray<number>, strides?: ReadonlyArray<number>,
       data?: Tensor.NumberType, channels?: number, width?: number, height?: number): TextureData {
     Logger.verbose('InferenceHandler', `Creating TextureData: shape:[${shape}], channels:${channels ? channels : 1}`);
-    const td = this.textureManager.createTexture(dataType, shape, strides, data, channels, width, height);
+    const td = this.textureHelper.createTexture(dataType, shape, strides, data, channels, width, height);
     return td;
   }
   createTextureDataFromLayout(layout: TextureLayout, dataType: Tensor.DataType, data?: Tensor.NumberType): TextureData {
     Logger.verbose('InferenceHandler', `Creating TextureData: layout:[${JSON.stringify(layout)}]`);
-    const td = this.textureManager.createTextureFromLayout(dataType, layout, data);
+    const td = this.textureHelper.createTextureFromLayout(dataType, layout, data);
     return td;
   }
   createBasicTextureLayout(

lib/backends/webgl/ops/reshape.ts

@@ -39,12 +39,13 @@ export class WebGLReshape extends Reshape {
       inputShapes.push(ShapeUtil.calculateReshapedDims(inputShape, inputs[1].integerData));
     }
     const inputTD = inferenceHandler.getOrCreate(inputs[0]);
+    const isInitializer = inferenceHandler.session.isInitializer(inputs[0]);
     const reshapedDims = this.getOutputShape(inferenceHandler, inputShapes);
     let packedShape = reshapedDims;
     if (inputTD.channels === 4) {
       packedShape = getPackedShape(reshapedDims);
     }
-    return [inferenceHandler.getTensor({
+    const newTD = {
       channels: inputTD.channels,
       dataType: inputs[0].type,
       texture: inputTD.texture,
@@ -54,7 +55,17 @@ export class WebGLReshape extends Reshape {
       strides: ShapeUtil.computeStrides(packedShape),
       unpackedShape: reshapedDims,
       arrayType: inputTD.arrayType
-    })];
+    };
+    const newTensor = new Tensor(newTD.unpackedShape, newTD.dataType, (id: Tensor.Id) => {
+      const values = inferenceHandler.textureHelper.readTexture(newTD, newTD.dataType, newTD.channels);
+      return values;
+    });
+    if (isInitializer) {
+      inferenceHandler.session.setTextureData(newTensor, newTD);
+    } else {
+      inferenceHandler.setTextureData(newTensor, newTD);
+    }
+    return [newTensor];
   }
   getPositionalFunction(inferenceHandler: WebGLInferenceHandler, inputShape: ReadonlyArray<number>, name?: string):
       GlslPositionalFunction {

lib/backends/webgl/program-manager.ts

@@ -62,10 +62,12 @@ export class ProgramManager {
   // tslint:disable-next-line:ban-types
   repo: Map<Object, Artifact>;  // this should be per-session object
   vertexShader: WebGLShader;
+  attributesBound: boolean;
 
   constructor(public profiler: Readonly<Profiler>, public glContext: WebGLContext) {
     this.repo = new Map();
     this.glContext = glContext;
+    this.attributesBound = false;
   }
   getArtifact(key: {}): Artifact|undefined {
     return this.repo.get(key);
@@ -84,7 +86,9 @@ export class ProgramManager {
       gl.useProgram(program);
       try {
         this.bindOutput(runData.outputTextureData);
-        this.bindAttributes(buildArtifact.attribLocations);
+        if (!this.attributesBound) {
+          this.bindAttributes(buildArtifact.attribLocations);
+        }
         this.bindUniforms(buildArtifact.uniformLocations, runData.uniformData);
         this.bindTextures(buildArtifact.uniformLocations, runData.inputTextureDatas);
       } catch (err) {
@@ -97,6 +101,7 @@ export class ProgramManager {
         } else {
           this.doDraw(buildArtifact, runData);
         }
+        gl.flush();
       });
       if (runData.postRun) {
         Logger.verbose('ProgramManager', 'PostRun');
@@ -138,7 +143,6 @@ export class ProgramManager {
       runData.draw(this.glContext, artifact);
     } else {
       this.glContext.draw();
-      this.glContext.gl.flush();
     }
   }
   protected doBlockDraw(artifact: Artifact, runData: RunData): void {
@@ -188,6 +192,7 @@ export class ProgramManager {
     const positionHandle = attribLocations.position.location as number;
     const textureCoordHandle = attribLocations.textureCoord.location as number;
     this.glContext.setVertexAttributes(positionHandle, textureCoordHandle);
+    this.attributesBound = true;
   }
   bindUniformArray(location: WebGLUniformLocation, type: string, value: number[]): void {
     const gl = this.glContext.gl;

lib/backends/webgl/session-handler.ts

@@ -29,20 +29,20 @@ import {WebGLTranspose} from './ops/transpose';
 import * as unaryOps from './ops/unary-op';
 import {ProgramManager} from './program-manager';
 import {TextureData} from './texture-data';
+import {TextureHelper} from './texture-helper';
 import {AlwaysKeepOriginalSizeStrategy, TextureLayoutStrategy} from './texture-layout-strategy';
-import {TextureManager} from './texture-manager';
 
 export class WebGLSessionHandler implements SessionHandler {
   programManager: ProgramManager;
-  textureManager: TextureManager;
+  textureHelper: TextureHelper;
   layoutStrategy: TextureLayoutStrategy;
   textureDataCache: Map<Tensor, TextureData>;
   initializers: Set<Tensor>;
 
   constructor(public readonly backend: WebGLBackend, public readonly context: Session.Context) {
     this.programManager = new ProgramManager(this.context.profiler, backend.glContext);
     this.layoutStrategy = new AlwaysKeepOriginalSizeStrategy(backend.glContext.maxTextureSize);
-    this.textureManager = new TextureManager(backend.glContext, this.layoutStrategy, this.context.profiler);
+    this.textureHelper = new TextureHelper(backend.glContext, this.layoutStrategy, this.context.profiler);
     this.textureDataCache = new Map();
   }
 
@@ -65,8 +65,9 @@ export class WebGLSessionHandler implements SessionHandler {
   }
   dispose(): void {
     this.programManager.dispose();
-    this.textureManager.dispose();
-    this.textureDataCache.forEach(td => this.textureManager.saveTexture(td.texture, [td.width, td.height]));
+    this.textureHelper.clearActiveTextures();
+    this.textureDataCache.forEach(td => this.textureHelper.releaseTexture(td.texture));
+    this.textureDataCache = new Map();
   }
   resolve(node: Graph.Node, domain: string, version: number): Operator {
     const op = this.createOperator(node, domain, version);

lib/backends/webgl/texture-manager.ts → lib/backends/webgl/texture-helper.ts

@@ -17,48 +17,37 @@ import {WebGLContext} from './webgl-context';
  *    Caching these is crucial to performance. These are In-use Textures
  *   2. textures which are not in use by any current ProgramInfo/Tensor
  *     These are called Free Textures
- * TextureManager is also used to help creating textures. For this it
+ * TextureHelper is also used to help creating textures. For this it
  * uses WebGLContext and TextureLayoutStrategy
  */
-export class TextureManager {
+export class TextureHelper {
   glContext: WebGLContext;
-  free: Map<string, WebGLTexture[]>;
   gl: WebGLRenderingContext;
   layoutStrategy: TextureLayoutStrategy;
   profiler: Readonly<Profiler>;
 
   constructor(context: WebGLContext, layoutStrategy: TextureLayoutStrategy, profiler: Readonly<Profiler>) {
     this.glContext = context;
     this.gl = context.gl;
-    this.free = new Map();
     this.layoutStrategy = layoutStrategy;
     this.profiler = profiler;
   }
   createTextureFromLayout(dataType: Tensor.DataType, layout: TextureLayout, data?: Tensor.NumberType) {
     let texture: WebGLTexture;
     const textureDataType = this.toEncoderType(dataType);
     const size = `${layout.width}-${layout.height}`;
-    const textureList = this.free.get(size);
-    if (!textureList || textureList.length === 0) {
-      Logger.verbose('TextureManager', `No cached texture; Creating new of size ${size}`);
-      texture = this.glContext.allocateTexture(
-          layout.width, layout.height, textureDataType, layout.channels, this.toTextureData(dataType, data));
-    } else {
-      Logger.verbose('TextureManager', `Found a texture in cache of size ${size}`);
-      texture = textureList.shift()!;
-      if (data) {
-        this.glContext.updateTexture(
-            texture, layout.width, layout.height, 'float' /*this.toEncoderType(dataType)*/, layout.channels,
-            this.toTextureData(dataType, data)!);
-      }
-    }
+
+    Logger.verbose('TextureHelper', `Creating new texture of size ${size}`);
+    texture = this.glContext.allocateTexture(
+        layout.width, layout.height, textureDataType, layout.channels, this.toTextureData(dataType, data));
+
     return {...layout, dataType, texture, arrayType: textureDataType};
   }
   createTexture(
       dataType: Tensor.DataType, shape: ReadonlyArray<number>, strides?: ReadonlyArray<number>,
       data?: Tensor.NumberType, channels?: number, width?: number, height?: number,
       unpackedShape?: ReadonlyArray<number>): TextureData {
-    return this.profiler.event('backend', 'TextureManager.createTexture', () => {
+    return this.profiler.event('backend', 'TextureHelper.createTexture', () => {
       if (!width || !height) {
         [width, height] = this.layoutStrategy.computeTextureWH(shape);
       }
@@ -80,23 +69,16 @@ export class TextureManager {
     if (!channels) {
       channels = 1;
     }
-    return this.profiler.event('backend', 'TextureManager.readTexture', () => {
+    return this.profiler.event('backend', 'TextureHelper.readTexture', () => {
       const dataSize = td.shape.reduce((a, b) => a * b) * channels!;
       const data = this.glContext.readTexture(
           td.texture, td.width, td.height, dataSize, this.toEncoderType(dataType), channels!);
       return this.toTensorData(dataType, data);
     });
   }
-  saveTexture(texture: WebGLTexture, dims: ReadonlyArray<number>): void {
-    return this.profiler.event('backend', 'TextureManager.saveTexture', () => {
-      const size = `${dims[0]}-${dims[1]}`;
-      Logger.verbose('TextureManager', `caching texture of size ${size}`);
-      let textureList = this.free.get(size);
-      if (!textureList) {
-        textureList = [];
-      }
-      textureList.push(texture);
-      this.free.set(size, textureList);
+  releaseTexture(texture: WebGLTexture): void {
+    return this.profiler.event('backend', 'TextureHelper.releaseTexture', () => {
+      this.glContext.deleteTexture(texture);
     });
   }
   createPaddedTexture(inputTextureData: TextureData, outputLayout: TextureLayout): TextureData {
@@ -179,7 +161,7 @@ export class TextureManager {
     //     throw new Error(`TensorData type ${dataType} is not supported`);
     // }
   }
-  dispose(): void {
-    this.free.forEach(value => value.forEach(t => this.glContext.deleteTexture(t)));
+  clearActiveTextures(): void {
+    this.glContext.clearActiveTextures();
   }
 }

lib/backends/webgl/webgl-context.ts

@@ -40,5 +40,6 @@ export interface WebGLContext {
   deleteTexture(texture: WebGLTexture): void;
   deleteProgram(program: WebGLProgram): void;
   getEncoder(dataType: Encoder.DataType, channels: number): DataEncoder;
+  clearActiveTextures(): void;
   dispose(): void;
 }

lib/execution-plan.ts

@@ -149,7 +149,8 @@ export class ExecutionPlan {
         thisValue.data;
         output.push(thisValue);
       });
-
+      Logger.verbose('ExecPlan', 'disposing of inferenceHandler');
+      inferenceHandler.dispose();
       return output;
     });
   }