add packed conv fuse

Author: xzhu1900

lib/backends/webgl/ops/concat_packed.ts

@@ -110,8 +110,9 @@ export class WebGLPackedConcat extends Concat implements WebGLOperator {
       `;
 
     return {
-      inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t)),
-      outputLayout: handler.createTextureLayoutFromShape(outputShape),
+      inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)),
+      outputLayout:
+          handler.createTextureLayoutFromShape(outputShape, 4, outputShape, {isPacked: true, reverseWH: true}),
       samplers,
       shaderSource,
       hasMain: true,
@@ -120,7 +121,7 @@ export class WebGLPackedConcat extends Concat implements WebGLOperator {
     };
   }
   createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
-    const inputTDs = inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i]));
+    const inputTDs = inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i], true));
     return {
       inputTextureDatas: inputTDs,
       outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),

lib/backends/webgl/ops/conv-pack.ts

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license.
 
+import {Attribute} from '../../../attribute';
 import {Logger} from '../../../instrument';
 import {Conv} from '../../../ops/conv';
 import {Tensor} from '../../../tensor';
@@ -36,6 +37,11 @@ export class WebGLConvPacked extends Conv {
     const outputShape = WebGLConv.calcOutputShape(xshape, kshape, this.dilations, this.pads, this.strides);
     const im2col = new WebGLIm2ColPacked(outputShape, kshape, this.dilations, this.pads, this.strides);
     const matmul = new WebGLMatMulPacked();
+    if (!!this.activation) {
+      const attributes = new Attribute(undefined);
+      attributes.set('__internal_activation', 'string', (this.activation));
+      matmul.initialize(attributes);
+    }
     const reshape = new WebGLReshapePacked();
     // shape for kernel reshape
     const shape =

lib/backends/webgl/ops/conv.ts

@@ -3,7 +3,7 @@
 
 import {Attribute} from '../../../attribute';
 import {Logger} from '../../../instrument';
-import {Conv, getActicationSnippet} from '../../../ops/conv';
+import {Conv} from '../../../ops/conv';
 import {Tensor} from '../../../tensor';
 import {PoolConvUtil} from '../../../util';
 import {getGlsl} from '../glsl-source';
@@ -12,6 +12,7 @@ import {Artifact, ProgramInfo, RunData, TextureLayout, WebGLOperator} from '../t
 import {WebGLContext} from '../webgl-context';
 
 import {WebGLConvPacked} from './conv-pack';
+import {getActicationSnippet} from './fuse_utils';
 
 export class WebGLConv extends Conv {
   unpackedGroupedConvImpl: WebGLUnpackedGroupedConv;

lib/backends/webgl/ops/fuse_utils.ts

@@ -0,0 +1,23 @@
+import {glslRelu, glslSigmoid} from './unary-op';
+
+export function getActicationSnippet(activation: string) {
+  let activationFunction = '';
+  let activationName = '';
+  switch (activation) {
+    case 'Relu':
+      activationName = glslRelu().name;
+      activationFunction = glslRelu().body;
+      break;
+    case 'Sigmoid':
+      activationName = glslSigmoid().name;
+      activationFunction = glslSigmoid().body;
+      break;
+    default:
+      activationName = '';
+      activationFunction = '';
+  }
+  const applyActivation = activation ? `
+  value = ${activationName}(value);` :
+                                       '';
+  return {activationFunction, applyActivation};
+}

lib/backends/webgl/ops/matmul-pack.ts

@@ -6,6 +6,7 @@ import {Tensor} from '../../../tensor';
 import {BroadcastUtil} from '../../../util';
 import {WebGLInferenceHandler} from '../inference-handler';
 import {ProgramInfo, RunData, WebGLOperator} from '../types';
+import {getActicationSnippet} from './fuse_utils';
 
 export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
   run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
@@ -25,8 +26,11 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
     const aRank = aShape.length;
     const bRank = bShape.length;
     const sharedDim = aShape[aShape.length - 1];
+
+    const {activationFunction, applyActivation} = getActicationSnippet(this.activation);
     // TODO:fix broadcasting
     const shaderSource = `
+      ${activationFunction}
       vec4 process(int indices[${rank}]) {
           int a[${aRank}];
           int b[${bRank}];
@@ -41,6 +45,7 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
               value += ${getA(aRank)}.ggaa * ${getB(bRank)}.baba;
           }
           ${processBias}
+          ${applyActivation}
           return value;
       }`;
     return {

lib/ops/conv.ts

@@ -3,7 +3,6 @@
 
 import {Attribute} from '../attribute';
 import {InferenceHandler} from '../backend';
-import {glslRelu, glslSigmoid} from '../backends/webgl/ops/unary-op';
 import {Operator} from '../operators';
 import {Tensor} from '../tensor';
 
@@ -92,25 +91,3 @@ export abstract class Conv implements Operator {
   protected strides: number[];
   protected activation: string;
 }
-
-export function getActicationSnippet(activation: string) {
-  let activationFunction = '';
-  let activationName = '';
-  switch (activation) {
-    case 'Relu':
-      activationName = glslRelu().name;
-      activationFunction = glslRelu().body;
-      break;
-    case 'Sigmoid':
-      activationName = glslSigmoid().name;
-      activationFunction = glslSigmoid().body;
-      break;
-    default:
-      activationName = '';
-      activationFunction = '';
-  }
-  const applyActivation = activation ? `
-  value = ${activationName}(value);` :
-                                       '';
-  return {activationFunction, applyActivation};
-}

lib/ops/matmul.ts

@@ -9,7 +9,9 @@ import {Tensor} from '../tensor';
 export abstract class MatMul implements Operator {
   abstract run(inferenceHandler: InferenceHandler, inputs: Tensor[]): Tensor[]|Promise<Tensor[]>;
 
-  initialize(attributes: Attribute): void {}
+  initialize(attributes: Attribute): void {
+    this.activation = attributes.getString('__internal_activation', '');
+  }
 
   checkInputs(inputs: Tensor[]): boolean {
     if (!inputs || inputs.length !== 2) {
@@ -38,4 +40,5 @@ export abstract class MatMul implements Operator {
 
     return true;
   }
+  protected activation: string;
 }