Revert "webgl: various optimizations (#13)" (#28)

This reverts commit 8acffd0.

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  --printMatches",
+    "test": "karma start --browsers ChromeTest --single-run",
     "test-debug": "karma start --browsers ChromeDebug  --printMatches",
-    "test-edge": "karma start --browsers Edge --single-run  --printMatches"
+    "test-edge": "karma start --browsers Edge --single-run"
   },
   "keywords": [
     "resnet",

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

@@ -229,13 +229,6 @@ 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,20 +20,18 @@ import {getPackedShape} from './utils';
  * Throughout WebGL backend operations TextureData is used as the data carrier
  */
 export class WebGLInferenceHandler implements InferenceHandler {
-  textureHelper: TextureHelper;
+  textureManager: TextureManager;
   programManager: ProgramManager;
   private tensorToTexture: Map<Tensor, TextureData>;
   private textureToTensor: Map<TextureData, Tensor>;
   constructor(public backend: WebGLBackend, public session: WebGLSessionHandler) {
-    this.textureHelper = session.textureHelper;
+    this.textureManager = session.textureManager;
     this.programManager = session.programManager;
     this.tensorToTexture = new Map();
     this.textureToTensor = new Map();
   }
   protected lookupTextureData(tensor: Tensor): TextureData|undefined {
-    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);
+    return this.session.isInitializer(tensor) ? this.session.getTextureData(tensor) : this.tensorToTexture.get(tensor);
   }
   getOrCreate(tensor: Tensor, layout?: TextureLayout): TextureData {
     let td = this.lookupTextureData(tensor);
@@ -66,7 +64,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.textureHelper.readTexture(td, td.dataType, td.channels);
+        const values = this.textureManager.readTexture(td, td.dataType, td.channels);
         return values;
       });
       this.setTextureData(tensor, td);
@@ -84,21 +82,20 @@ export class WebGLInferenceHandler implements InferenceHandler {
         channels === 1 ? tensor.dims.slice() : getPackedShape(tensor.dims.slice()), channels, unpackedShape);
   }
   dispose(): void {
-    this.textureHelper.clearActiveTextures();
-    this.tensorToTexture.forEach(td => this.textureHelper.releaseTexture(td.texture));
+    this.tensorToTexture.forEach(td => this.textureManager.saveTexture(td.texture, [td.width, td.height]));
     this.tensorToTexture = new Map();
     this.textureToTensor = new Map();
   }
   createTextureData(
       dataType: Tensor.DataType, shape: number[], strides?: 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.textureHelper.createTexture(dataType, shape, strides, data, channels, width, height);
+    const td = this.textureManager.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.textureHelper.createTextureFromLayout(dataType, layout, data);
+    const td = this.textureManager.createTextureFromLayout(dataType, layout, data);
     return td;
   }
   createBasicTextureLayout(shape: number[], channels = 1, unpackedShape?: number[], prefs?: WidthHeightPrefs):

lib/backends/webgl/ops/conv.ts

@@ -213,8 +213,7 @@ export class WebGLConv extends Conv {
       inputLayouts: inputs.length === 3 ? [im2colLayout, kLayout, bLayout!] : [im2colLayout, kLayout],
       outputLayout,
       shaderSource,
-      params: {'sharedDim': sharedDim, 'sharedDimReadSize': sharedDimReadSize},
-      blockSize: this.calcBlockSize(outputLayout)
+      params: {'sharedDim': sharedDim, 'sharedDimReadSize': sharedDimReadSize}
     };
   }
   createDotProdRunData(inferenceHandler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
@@ -289,15 +288,15 @@ export class WebGLConv extends Conv {
     return outputShape;
   }
   protected calcSharedDimReadSize(sharedDim: number): number {
-    const preferredBatchSize = 8;
+    const preferredBatchSize = 16;
     if (sharedDim < preferredBatchSize || sharedDim % preferredBatchSize !== 0) {
       return sharedDim;
     }
     return preferredBatchSize;
   }
   protected calcBlockSize(outputLayout: TextureLayout): [number, number]|undefined {
-    const preferredRowCount = 32;
-    const preferredColCount = 32;
+    const preferredRowCount = 64;
+    const preferredColCount = 64;
     if (outputLayout.height < preferredRowCount) {
       return undefined;
     }

lib/backends/webgl/ops/reshape.ts

@@ -38,13 +38,12 @@ 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);
     }
-    const newTD = {
+    return [inferenceHandler.getTensor({
       channels: inputTD.channels,
       dataType: inputs[0].type,
       texture: inputTD.texture,
@@ -54,17 +53,7 @@ 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: number[], name?: string):
       GlslPositionalFunction {

lib/backends/webgl/program-manager.ts

@@ -62,12 +62,10 @@ 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);
@@ -86,9 +84,7 @@ export class ProgramManager {
       gl.useProgram(program);
       try {
         this.bindOutput(runData.outputTextureData);
-        if (!this.attributesBound) {
-          this.bindAttributes(buildArtifact.attribLocations);
-        }
+        this.bindAttributes(buildArtifact.attribLocations);
         this.bindUniforms(buildArtifact.uniformLocations, runData.uniformData);
         this.bindTextures(buildArtifact.uniformLocations, runData.inputTextureDatas);
       } catch (err) {
@@ -101,7 +97,6 @@ export class ProgramManager {
         } else {
           this.doDraw(buildArtifact, runData);
         }
-        gl.flush();
       });
       if (runData.postRun) {
         Logger.verbose('ProgramManager', 'PostRun');
@@ -143,6 +138,7 @@ export class ProgramManager {
       runData.draw(this.glContext, artifact);
     } else {
       this.glContext.draw();
+      this.glContext.gl.flush();
     }
   }
   protected doBlockDraw(artifact: Artifact, runData: RunData): void {
@@ -192,7 +188,6 @@ 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;
-  textureHelper: TextureHelper;
+  textureManager: TextureManager;
   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.textureHelper = new TextureHelper(backend.glContext, this.layoutStrategy, this.context.profiler);
+    this.textureManager = new TextureManager(backend.glContext, this.layoutStrategy, this.context.profiler);
     this.textureDataCache = new Map();
   }
 
@@ -65,9 +65,8 @@ export class WebGLSessionHandler implements SessionHandler {
   }
   dispose(): void {
     this.programManager.dispose();
-    this.textureHelper.clearActiveTextures();
-    this.textureDataCache.forEach(td => this.textureHelper.releaseTexture(td.texture));
-    this.textureDataCache = new Map();
+    this.textureManager.dispose();
+    this.textureDataCache.forEach(td => this.textureManager.saveTexture(td.texture, [td.width, td.height]));
   }
   resolve(node: Graph.Node, domain: string, version: number): Operator {
     const op = this.createOperator(node, domain, version);

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

@@ -17,36 +17,47 @@ 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
- * TextureHelper is also used to help creating textures. For this it
+ * TextureManager is also used to help creating textures. For this it
  * uses WebGLContext and TextureLayoutStrategy
  */
-export class TextureHelper {
+export class TextureManager {
   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}`;
-
-    Logger.verbose('TextureHelper', `Creating new texture of size ${size}`);
-    texture = this.glContext.allocateTexture(
-        layout.width, layout.height, textureDataType, layout.channels, this.toTextureData(dataType, data));
-
+    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)!);
+      }
+    }
     return {...layout, dataType, texture, arrayType: textureDataType};
   }
   createTexture(
       dataType: Tensor.DataType, shape: number[], strides?: number[], data?: Tensor.NumberType, channels?: number,
       width?: number, height?: number, unpackedShape?: number[]): TextureData {
-    return this.profiler.event('backend', 'TextureHelper.createTexture', () => {
+    return this.profiler.event('backend', 'TextureManager.createTexture', () => {
       if (!width || !height) {
         [width, height] = this.layoutStrategy.computeTextureWH(shape);
       }
@@ -68,16 +79,23 @@ export class TextureHelper {
     if (!channels) {
       channels = 1;
     }
-    return this.profiler.event('backend', 'TextureHelper.readTexture', () => {
+    return this.profiler.event('backend', 'TextureManager.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);
     });
   }
-  releaseTexture(texture: WebGLTexture): void {
-    return this.profiler.event('backend', 'TextureHelper.releaseTexture', () => {
-      this.glContext.deleteTexture(texture);
+  saveTexture(texture: WebGLTexture, dims: 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);
     });
   }
   createPaddedTexture(inputTextureData: TextureData, outputLayout: TextureLayout): TextureData {
@@ -160,7 +178,7 @@ export class TextureHelper {
     //     throw new Error(`TensorData type ${dataType} is not supported`);
     // }
   }
-  clearActiveTextures(): void {
-    this.glContext.clearActiveTextures();
+  dispose(): void {
+    this.free.forEach(value => value.forEach(t => this.glContext.deleteTexture(t)));
   }
 }

lib/backends/webgl/webgl-context.ts

@@ -40,6 +40,5 @@ 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,8 +149,7 @@ export class ExecutionPlan {
         thisValue.data;
         output.push(thisValue);
       });
-      Logger.verbose('ExecPlan', 'disposing of inferenceHandler');
-      inferenceHandler.dispose();
+
       return output;
     });
   }