add texture cache support for pack/unpack mode (#290)

* squash change

* uncomment test verification

* clean up

* clean up and remove unused files

* adding comment

* remove unused code

* fix reshape merge

* add guid as tensor id

* add file

* fix test failure

Author: xzhu1900

lib/backends/webgl/glsl-coordinate-lib.ts

@@ -690,7 +690,7 @@ export class CoordsGlslLib extends GlslLib {
           return ${texFuncSnippet}(${unpackedCoordsSnippet});
         }
       `;
-    return new GlslLibRoutine(source);
+    return new GlslLibRoutine(source, ['coordinates.getOutputCoords']);
   }
 
   /**
@@ -1216,4 +1216,27 @@ export class CoordsGlslLib extends GlslLib {
         }
         `;
   }
+
+  /**
+   * Produces a packed value getter function for the name and rank given
+   * If a transpose is set proper offsetToCoords mapping will be used
+   * @param name name of the function
+   * @param rank rank of the input
+   * @param transpose whether or not should generate a transpose variation
+   */
+  protected getPackedValueFrom(varName: string, rank: number, width: number, height: number, transpose: boolean):
+      string {
+    let name = `_${varName}_Pack`;
+    if (transpose) {
+      name = name + '_T';
+    }
+    const glsl = getGlsl(this.context.glContext.version);
+    return `
+        vec4 ${name}(int m[${rank}]) {
+          int offset = indicesToOffset_${varName}(m);
+          vec2 coords = offsetToCoords(offset, ${width}, ${height});
+          return ${glsl.texture2D}(${varName}, coords);
+        }
+        `;
+  }
 }

lib/backends/webgl/inference-handler.ts

@@ -16,9 +16,16 @@ import {Artifact, RunData, TextureData, TextureLayout, WebGLOperator} from './ty
 import {getPackedShape} from './utils';
 
 export class WebGLInferenceHandler implements InferenceHandler {
-  private textureDataCache: Map<Tensor.Id, TextureData>;
+  private packedTextureDataCache: Map<Tensor.Id, TextureData>;
+  private unpackedTextureDataCache: Map<Tensor.Id, TextureData>;
+  private pack2unpackMap: Map<Tensor.Id, Tensor.Id>;
+  private unpack2packMap: Map<Tensor.Id, Tensor.Id>;
   constructor(public session: WebGLSessionHandler) {
-    this.textureDataCache = new Map();
+    this.packedTextureDataCache = new Map();
+    this.unpackedTextureDataCache = new Map();
+
+    this.pack2unpackMap = new Map();
+    this.unpack2packMap = new Map();
   }
 
   run(op: WebGLOperator, inputs: Tensor[]): Tensor[] {
@@ -33,36 +40,18 @@ export class WebGLInferenceHandler implements InferenceHandler {
     return [runData.outputTextureData.tensor];
   }
 
-  /**
-   * Check the runData's input texture mode with the program's artifact.
-   * If the artifact expects a packed input, while the RunData's input
-   * is unpacked, perform a pack operation on this input to align the
-   * texture mode with artifact. Similar on unpacked input.
-   */
   checkAndUpdateTextureForm(artifact: Artifact, runData: RunData) {
     // pack/unpack inputs
-    runData.inputTextureDatas.forEach(input => {
+    for (let i = 0; i < runData.inputTextureDatas.length; ++i) {
+      const input = runData.inputTextureDatas[i];
       if (input.isPacked && !artifact.programInfo.expectPackedInputs) {
-        // unpack this input
-        const unpacked = this.unpack(input);
-        input.height = unpacked.height;
-        input.isPacked = unpacked.isPacked;
-        input.texture = unpacked.texture;
-        input.width = unpacked.width;
-
+        runData.inputTextureDatas[i] = this.unpack(input);
       } else if (!input.isPacked && artifact.programInfo.expectPackedInputs) {
-        // pack this input
-        const packed = this.pack(input);
-        input.height = packed.height;
-        input.isPacked = packed.isPacked;
-        input.texture = packed.texture;
-        input.width = packed.width;
+        runData.inputTextureDatas[i] = this.pack(input);
       }
-    });
+    }
   }
   runProgram(artifact: Artifact, runData: RunData) {
-    // if the runData has different expected texture pack/unpack mode, process pack/unpack
-    // operation on the texture before executing the kernel.
     this.checkAndUpdateTextureForm(artifact, runData);
 
     // output should match
@@ -84,15 +73,28 @@ export class WebGLInferenceHandler implements InferenceHandler {
    *   Creates a texture data object associated with the given tensor.
    * @param tensor the tensor with data to upload
    */
-  getOrCreateTextureData(tensor: Tensor, layout?: TextureLayout) {
-    let td = this.getTextureData(tensor.dataId);
+  getOrCreateTextureData(tensor: Tensor, layout?: TextureLayout, isPacked = false) {
+    let td = this.getTextureData(tensor.dataId, isPacked);
     if (!td) {
       Logger.verbose('InferenceHandler', `Creating new TextureData for dims: [${tensor.dims}]`);
       if (!layout) {
         layout = this.createTextureLayoutFromShape(tensor.dims.slice());
       }
-      // graph inputs or initializers
-      td = this.createTextureData(layout, tensor.type, tensor.numberData, tensor, Encoder.Usage.UploadOnly);
+      // if we don't find the texture data with specific pack mode in the cache, try with the different
+      // pack mode to see if the tensor is cached using that pack mode. If succeed, we can return this
+      // tensor data and later apply a pack/unpack op on this texture, no need to create a new one here.
+      td = this.getTextureData(tensor.dataId, !isPacked);
+      if (!td) {
+        if (isPacked) {
+          const unpackedTextureLayout = this.getOrCreateTextureLayout(tensor, 1, false, [], true);
+          const unpackedTextureData = this.createTextureData(
+              unpackedTextureLayout, tensor.type, tensor.numberData, tensor, Encoder.Usage.UploadOnly);
+          td = this.pack(unpackedTextureData);
+        } else {
+          td = this.createTextureData(
+              layout, tensor.type, tensor.numberData, tensor, Encoder.Usage.UploadOnly, isPacked);
+        }
+      }
     } else {
       Logger.verbose('InferenceHandler', `Retrieving TextureData from cache: [${tensor.dims}]`);
     }
@@ -104,7 +106,7 @@ export class WebGLInferenceHandler implements InferenceHandler {
    * Usage = Encoder.Usage.Default.
    * @param dataType the tensor data type
    */
-  createTextureDataFromLayout(layout: TextureLayout, dataType: Tensor.DataType): TextureData {
+  createTextureDataFromLayout(layout: TextureLayout, dataType: Tensor.DataType, isPacked = false): TextureData {
     return this.createTextureData(layout, dataType);
   }
 
@@ -118,13 +120,14 @@ export class WebGLInferenceHandler implements InferenceHandler {
    * @param tensor the tensor to bind. tensor's data is ignored.
    */
   createTextureDataFromLayoutBindTensor(
-      layout: TextureLayout, dataType: Tensor.DataType, data: Tensor.NumberType, tensor: Tensor): TextureData {
-    return this.createTextureData(layout, dataType, data, tensor, Encoder.Usage.UploadOnly);
+      layout: TextureLayout, dataType: Tensor.DataType, data: Tensor.NumberType, tensor: Tensor,
+      isPacked = false): TextureData {
+    return this.createTextureData(layout, dataType, data, tensor, Encoder.Usage.UploadOnly, isPacked);
   }
 
   private createTextureData(
       layout: TextureLayout, dataType: Tensor.DataType, data?: Tensor.NumberType, tensor?: Tensor,
-      usage?: Encoder.Usage): TextureData {
+      usage?: Encoder.Usage, isPacked = false): TextureData {
     Logger.verbose('InferenceHandler', `Creating TextureData: layout:[${JSON.stringify(layout)}]`);
     const texture = this.session.textureManager.createTextureFromLayout(dataType, layout, data, usage);
     return this.createTextureDataFromTexture(layout, dataType, texture, tensor);
@@ -137,8 +140,9 @@ export class WebGLInferenceHandler implements InferenceHandler {
    * @param texture the WebGLTexture object to share
    * @param tensorId the tensor ID of the shared tensor data
    */
-  createSharedTextureData(layout: TextureLayout, dataType: Tensor.DataType, texture: WebGLTexture, tensorId: Tensor.Id):
-      TextureData {
+  createSharedTextureData(
+      layout: TextureLayout, dataType: Tensor.DataType, texture: WebGLTexture, tensorId?: Tensor.Id,
+      isPacked = false): TextureData {
     return this.createTextureDataFromTexture(layout, dataType, texture, undefined, tensorId);
   }
 
@@ -155,29 +159,32 @@ export class WebGLInferenceHandler implements InferenceHandler {
                   undefined, undefined, tensorId),
       texture
     };
-    this.setTextureData(textureData.tensor.dataId, textureData);
+    this.setTextureData(textureData.tensor.dataId, textureData, layout.isPacked);
     return textureData;
   }
 
-  getTextureData(tensorId: Tensor.Id): TextureData|undefined {
-    return this.session.isInitializer(tensorId) ? this.session.getTextureData(tensorId) :
-                                                  this.textureDataCache.get(tensorId);
+  getTextureData(tensorId: Tensor.Id, isPacked = false): TextureData|undefined {
+    return this.session.isInitializer(tensorId) ?
+        this.session.getTextureData(tensorId, isPacked) :
+        isPacked ? this.packedTextureDataCache.get(tensorId) : this.unpackedTextureDataCache.get(tensorId);
   }
-  setTextureData(tensorId: Tensor.Id, td: TextureData): void {
+  setTextureData(tensorId: Tensor.Id, td: TextureData, isPacked = false): void {
     if (this.session.isInitializer(tensorId)) {
-      this.session.setTextureData(tensorId, td);
+      this.session.setTextureData(tensorId, td, isPacked);
     } else {
-      this.textureDataCache.set(tensorId, td);
+      isPacked ? this.packedTextureDataCache.set(tensorId, td) : this.unpackedTextureDataCache.set(tensorId, td);
     }
   }
-
+  isTextureLayoutCached(tensor: Tensor, isPacked = false): boolean {
+    return !!this.getTextureData(tensor.dataId, isPacked);
+  }
   /**
    * Create a TextureLayout object from a tensor. If a related texture data is found, returns the cached texture layout.
    */
   getOrCreateTextureLayout(
       tensor: Tensor, channels: 1|4 = 1, isPacked = false, unpackedShape?: ReadonlyArray<number>,
       reverseWH = false): TextureLayout {
-    const td = this.getTextureData(tensor.dataId);
+    const td = this.getTextureData(tensor.dataId, isPacked);
     if (td) {
       return td;
     }
@@ -229,14 +236,17 @@ export class WebGLInferenceHandler implements InferenceHandler {
       isPacked,
       shape: inferredDims,
       strides: ShapeUtil.computeStrides(inferredDims),
-      unpackedShape
+      unpackedShape,
+      reversedWH: (prefs && prefs.reverseWH)
     };
   }
 
   dispose(): void {
     this.session.textureManager.clearActiveTextures();
-    this.textureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
-    this.textureDataCache = new Map();
+    this.packedTextureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
+    this.packedTextureDataCache = new Map();
+    this.unpackedTextureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
+    this.unpackedTextureDataCache = new Map();
   }
 
   readTexture(textureData: TextureData): Tensor.NumberType {
@@ -252,6 +262,10 @@ export class WebGLInferenceHandler implements InferenceHandler {
   }
 
   pack(input: TextureData): TextureData {
+    const cachedId = this.unpack2packMap.get(input.tensor.dataId);
+    if (cachedId) {
+      return this.packedTextureDataCache.get(cachedId)!;
+    }
     const key = `${input.shape}`;
     let op = this.session.packOpCache.get(key);
     if (!op) {
@@ -266,10 +280,15 @@ export class WebGLInferenceHandler implements InferenceHandler {
     }
     const runData = op.createRunData(this, artifact.programInfo, [input.tensor]);
     this.runProgram(artifact, runData);
+    this.unpack2packMap.set(input.tensor.dataId, runData.outputTextureData.tensor.dataId);
     return runData.outputTextureData;
   }
 
   unpack(input: TextureData): TextureData {
+    const cachedId = this.pack2unpackMap.get(input.tensor.dataId);
+    if (cachedId) {
+      return this.unpackedTextureDataCache.get(cachedId)!;
+    }
     // For unpacked kernel, cache it by using input's unpackedShape as cache key.
     // Note that we need to use input.unpackedShape instead of input.shape here,
     // as the shape infers the packed texture shape. Different unpackedShape can have the
@@ -290,6 +309,7 @@ export class WebGLInferenceHandler implements InferenceHandler {
     }
     const runData = op.createRunData(this, artifact.programInfo, [input.tensor]);
     this.runProgram(artifact, runData);
+    this.pack2unpackMap.set(input.tensor.dataId, runData.outputTextureData.tensor.dataId);
     return runData.outputTextureData;
   }
 }

lib/backends/webgl/ops/binary-op.ts

@@ -22,6 +22,10 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
     const inputLayouts = handler.session.pack ?
         inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)) :
         inputs.map(t => handler.getOrCreateTextureLayout(t));
+    const ouputLayout = handler.session.pack ?
+        handler.createTextureLayoutFromShape(inputs[0].dims, 4, inputs[0].dims, {isPacked: true, reverseWH: true}) :
+        handler.createTextureLayoutFromShape(inputs[0].dims);
+
     const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
     if (isBroadcast) {
       const outputShape = BroadcastUtil.calcShape(inputs[0].dims, inputs[1].dims, false);
@@ -33,6 +37,8 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
       const bRank = inputs[1].dims.length !== 0 ? inputs[1].dims.length : 1;
       const aBcast = inputs[0].dims.length !== 0 ? `bcastIndices_A(indices, aindices);` : `aindices[0] = 0;`;
       const bBcast = inputs[1].dims.length !== 0 ? `bcastIndices_B(indices, bindices);` : `bindices[0] = 0;`;
+
+      // TODO: for packed tensors, we need to implement logic to caculate textCoords for broadcast tensor
       const shaderSource = `
       ${this.glslFunc.body}
       float process(int indices[${outputRank}]) {
@@ -51,6 +57,8 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
         outputLayout,
         samplers: ['A', 'B'],
         shaderSource,
+        expectPackedInputs: handler.session.pack,
+        expectPackedOutputs: handler.session.pack
       };
     }
     const glsl = getGlsl(handler.session.backend.glContext.version);
@@ -67,7 +75,7 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
       return {
         hasMain: true,
         inputLayouts,
-        outputLayout: handler.createTextureLayoutFromShape(inputs[0].dims),
+        outputLayout: ouputLayout,
         samplers: ['A', 'B'],
         shaderSource,
         expectPackedInputs: true,

lib/backends/webgl/ops/reshape-packed.ts

@@ -117,10 +117,17 @@ export class WebGLReshapePacked extends Reshape implements WebGLOperator {
   createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
     const inputTDs =
         [handler.getOrCreateTextureData(inputs[0], handler.getOrCreateTextureLayout(inputs[0], 1, false, [], false))];
+    let outputLayout = this.originalOutputLayout;
+    if (outputLayout === undefined) {
+      const originInputShape = inputs[0].dims;
+      const outputShape = ShapeUtil.calculateReshapedDims(originInputShape, inputs[1].integerData);
+      outputLayout =
+          handler.createTextureLayoutFromShape(outputShape, 4, outputShape, {isPacked: true, reverseWH: true});
+    }
     // return run data for reshape. Here, we use the original calculate outputLayout to create the real output layout.
     return {
       inputTextureDatas: inputTDs,
-      outputTextureData: handler.createTextureDataFromLayout(this.originalOutputLayout, inputTDs[0].tensor.type),
+      outputTextureData: handler.createTextureDataFromLayout(outputLayout, inputTDs[0].tensor.type),
       uniformData: {}
     };
   }

lib/backends/webgl/ops/reshape.ts

@@ -43,6 +43,6 @@ export function reshape(
     unpackedShape: reshapedDims,
   };
 
-  const newTextureData = inferenceHandler.createSharedTextureData(newTextureLayout, input.type, inputTD.texture, {});
+  const newTextureData = inferenceHandler.createSharedTextureData(newTextureLayout, input.type, inputTD.texture);
   return newTextureData.tensor;
 }

lib/backends/webgl/ops/uint8-encode.ts

@@ -77,7 +77,7 @@ export class WebGLUint8Encode {
     const encoder = inferenceHandler.session.backend.glContext.getEncoder('byte', 4);
     const texture =
         inferenceHandler.session.backend.glContext.allocateTexture(outputLayout.width, outputLayout.height, encoder);
-    const outputTextureData = inferenceHandler.createSharedTextureData(outputLayout, 'uint8', texture, {});
+    const outputTextureData = inferenceHandler.createSharedTextureData(outputLayout, 'uint8', texture);
     const runData = {inputTextureDatas: [input], outputTextureData, uniformData: {}};
 
     inferenceHandler.session.programManager.run(artifact, runData);

lib/backends/webgl/ops/unpack.ts

@@ -6,7 +6,6 @@ import {getGlsl} from '../glsl-source';
 import {WebGLInferenceHandler} from '../inference-handler';
 import {ProgramInfo, RunData, WebGLOperator} from '../types';
 import {getCoordsDataType} from '../utils';
-
 import {getChannels, unpackFromChannel} from './packing_utils';
 
 export class WebGLUnpack implements WebGLOperator {
@@ -18,7 +17,7 @@ export class WebGLUnpack implements WebGLOperator {
       throw new Error(`Pack kernel should have input tensor count to 1.`);
     }
 
-    const inputTexture = handler.getTextureData(inputs[0].dataId);
+    const inputTexture = handler.getTextureData(inputs[0].dataId, true);
     if (!inputTexture) {
       throw new Error(`packed input texture must exist`);
     }
@@ -49,7 +48,7 @@ export class WebGLUnpack implements WebGLOperator {
       `;
 
     return {
-      inputLayouts: [handler.getOrCreateTextureLayout(inputs[0])],
+      inputLayouts: [handler.getOrCreateTextureLayout(inputs[0], 4, true, inputs[0].dims, true)],
       outputLayout,
       samplers: ['A'],
       shaderSource,
@@ -59,7 +58,7 @@ export class WebGLUnpack implements WebGLOperator {
     };
   }
   createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
-    const inputTDs = [handler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0])];
+    const inputTDs = [handler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0], true)];
     return {
       inputTextureDatas: inputTDs,
       outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),

lib/backends/webgl/program-manager.ts

@@ -105,11 +105,12 @@ ${fragShaderScript}
     return program;
   }
   bindOutput(td: TextureData): void {
+    const width = td.width;
+    const height = td.height;
     Logger.verbose(
         'ProrgramManager',
-        `Binding output texture to Framebuffer: w/h=${td.width}/${td.height}, shape=${td.shape}, type=${
-            td.tensor.type}`);
-    this.glContext.attachFramebuffer(td.texture, td.width, td.height);
+        `Binding output texture to Framebuffer: w/h=${width}/${height}, shape=${td.shape}, type=${td.tensor.type}`);
+    this.glContext.attachFramebuffer(td.texture, width, height);
   }
   bindAttributes(attribLocations: Artifact.AttribLocations): void {
     const positionHandle = attribLocations.position;