Adding Conv+Clip fusion (#8102)

js/web/lib/onnxjs/backends/webgl/ops/clip.ts

@@ -1,8 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+/* eslint-disable no-invalid-this */
+
 import {Clip} from '../../../ops/clip';
 import {Tensor} from '../../../tensor';
+import {FunctionType} from '../glsl-definitions';
 import {getGlsl} from '../glsl-source';
 import {WebGLInferenceHandler} from '../inference-handler';
 import {ProgramInfo, RunData, WebGLOperator} from '../types';
@@ -39,3 +42,16 @@ export class WebGLClip extends Clip implements WebGLOperator {
     };
   }
 }
+// used for fusion
+export function glslClip() {
+  const name = 'clip_';
+  const body = `
+    float ${name}(float a, float max, float min) {
+      return clamp(a, min, max);
+    }
+    vec4 ${name}(vec4 v, float max, float min) {
+      return clamp(v, min, max);
+    }
+    `;
+  return {body, name, type: FunctionType.ValueBased};
+}

js/web/lib/onnxjs/backends/webgl/ops/conv-pack.ts

@@ -49,6 +49,10 @@ export class WebGLConvPacked extends Conv {
     if (this.activation) {
       const attributes = new Attribute(undefined);
       attributes.set('__internal_activation', 'string', (this.activation));
+      if (this.activation === 'Clip') {
+        attributes.set('__clip_max', 'float', this.clipMax);
+        attributes.set('__clip_min', 'float', this.clipMin);
+      }
       this.matmul.initialize(attributes);
     }
     // shape for kernel reshape

js/web/lib/onnxjs/backends/webgl/ops/conv.ts

@@ -35,9 +35,10 @@ export class WebGLConv extends Conv {
 
   run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
     const packMode = inferenceHandler.session.pack;
+    const isPointwise = this.kernelShape[0] === 1 && this.kernelShape[1] === 1;
     if (this.group > 1) {
       return this.unpackedGroupedConvImpl.run(inferenceHandler, inputs);
-    } else if (packMode && inputs[0].dims.length === 4 && inputs[0].dims[0] === 1) {
+    } else if (packMode && inputs[0].dims.length === 4 && inputs[0].dims[0] === 1 && !isPointwise) {
       return this.packedConvImpl.run(inferenceHandler, inputs);
     } else {
       return this.unpackedConvImpl.run(inferenceHandler, inputs);
@@ -86,12 +87,16 @@ export class WebGLUnpackedGroupedConv extends Conv implements WebGLOperator {
             this.kernelShape}, pads:${this.pads}, strides:${this.strides}`);
     const outputShape = WebGLConv.calcOutputShape(xShape, wShape, this.dilations, this.pads, this.strides);
     const glsl = getGlsl(handler.session.backend.glContext.version);
-
+    const additionalVars = this.activation === 'Clip' ? `
+    const float min = float(${this.clipMin});
+    const float max = float(${this.clipMax});` :
+                                                        '';
     const {activationFunction, applyActivation} = getActicationSnippet(this.activation);
 
     const shaderSource = `
     const ivec2 strides = ivec2(${this.strides[0]}, ${this.strides[1]});
     const ivec2 pads = ivec2(${this.pads[0]}, ${this.pads[1]});
+    ${additionalVars}
     ${activationFunction}
     void main() {
       ivec4 coords = getOutputCoords();
@@ -334,11 +339,14 @@ export class WebGLUnpackedConv extends Conv {
     if (inputs.length === 3) {
       samplers.push('B');
     }
-
+    const additionalVars = this.activation === 'Clip' ? `
+    const float min = float(${this.clipMin});
+    const float max = float(${this.clipMax});` :
+                                                        '';
     const {activationFunction, applyActivation} = getActicationSnippet(this.activation);
-
     const glsl = getGlsl(inferenceHandler.session.backend.glContext.version);
     const shaderSource = `
+    ${additionalVars}
     ${activationFunction}
     float process(int indices[${rank}]) {
       int b[1];

js/web/lib/onnxjs/backends/webgl/ops/fuse-utils.ts

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+import {glslClip} from './clip';
 import {glslRelu, glslSigmoid} from './unary-op';
 
 export function getActicationSnippet(activation: string) {
@@ -15,13 +16,20 @@ export function getActicationSnippet(activation: string) {
       activationName = glslSigmoid().name;
       activationFunction = glslSigmoid().body;
       break;
+    case 'Clip':
+      activationName = glslClip().name;
+      activationFunction = glslClip().body;
+      break;
     default:
       // TODO: adding other activations that can be fused.
       activationName = '';
       activationFunction = '';
   }
-  const applyActivation = activation ? `
-  value = ${activationName}(value);` :
+
+  const applyActivation = activation ? (activation === 'Clip' ? `
+  value = ${activationName}(value, max, min);` :
+                                                                `
+  value = ${activationName}(value);`) :
                                        '';
   return {activationFunction, applyActivation};
 }

js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts

@@ -17,7 +17,7 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
   }
   createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo {
     const hasBias = inputs.length > 2;
-    const processBias = hasBias ? 'result += getBiasAtOutCoords();' : '';
+    const processBias = hasBias ? 'value += getBiasAtOutCoords();' : '';
     const aShape = inputs[0].dims;
     const bShape = inputs[1].dims;
     const outputShape = BroadcastUtil.calcShape(aShape, bShape, true);
@@ -34,9 +34,13 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
     const coordsDataType = getCoordsDataType(outputShape.length);
     const outRank = outputShape.length;
     const allGlChannels = ['x', 'y', 'z', 'w', 'u', 'v'];
-
+    const additionalVars = this.activation === 'Clip' ? `
+    float min = float(${this.clipMin});
+    float max = float(${this.clipMax});` :
+                                                        '';
     const {activationFunction, applyActivation} = getActicationSnippet(this.activation);
     const shaderSource = `
+      ${additionalVars}
       ${activationFunction}
       void main() {
         ${coordsDataType} rc = getOutputCoords();
@@ -44,16 +48,16 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
         rc.${allGlChannels[outRank - 1]} = rc.${allGlChannels[outRank - 2]};
         rc.${allGlChannels[outRank - 2]} = lastDim;
 
-        vec4 result = vec4(0);
+        vec4 value = vec4(0);
         for (int i = 0; i < ${sharedDimIndex}; i++) {
           vec4 a = getA(${getA(allGlChannels, aRank)});
           vec4 b = getB(${getB(allGlChannels, bRank)});
-          result += (a.rrbb * b.rgrg);
-          result += (a.ggaa * b.baba);
+          value += (a.rrbb * b.rgrg);
+          value += (a.ggaa * b.baba);
         }
         ${processBias}
         ${applyActivation}
-        ${glsl.output} = result;
+        ${glsl.output} = value;
       }`;
     return {
       name: 'WebGLMatMulPacked',

js/web/lib/onnxjs/graph.ts

@@ -53,6 +53,7 @@ export declare namespace Graph {
   export interface Transformer {
     removeAllIdentityNodes(): void;
     removeAllDropoutNodes(): void;
+    fuseConvActivationNodes(): void;
     // TODO: add generic functions to manipulate the graph
   }
 
@@ -560,7 +561,7 @@ class GraphImpl implements Graph, Graph.Transformer {
     // apply common transform
     this.removeAllIdentityNodes();
     this.removeAllDropoutNodes();
-
+    this.fuseConvActivationNodes();
     // apply initializer specific transform
     if (graphInitializer) {
       graphInitializer.transformGraph(this);
@@ -743,6 +744,7 @@ class GraphImpl implements Graph, Graph.Transformer {
       // TODO: add other activation methods
       case 'Relu':
       case 'Sigmoid':
+      case 'Clip':
         return true;
       default:
         return false;
@@ -754,10 +756,16 @@ class GraphImpl implements Graph, Graph.Transformer {
       if (node.opType === 'Conv') {
         const next = this._allData[node.outputs[0]]._to;
         if (next.length === 1 && this.isActivation(this._nodes[next[0]])) {
-          node.attributes.set('__internal_activation', 'string', (this._nodes[next[0]].opType));
+          const child = this._nodes[next[0]];
+          node.attributes.set('__internal_activation', 'string', (child.opType));
+          // TODO: need add support for Clip after opset 11, which has min/max as inputs
+          if (child.opType === 'Clip') {
+            node.attributes.set('__clip_min', 'float', child.attributes.getFloat('min'));
+            node.attributes.set('__clip_max', 'float', child.attributes.getFloat('max'));
+          }
           this.deleteNode(next[0]);
         }
       }
     }
   }
-}
+}

js/web/lib/onnxjs/ops/conv.ts

@@ -18,6 +18,8 @@ export abstract class Conv implements Operator {
     this.pads = attributes.getInts('pads', [0, 0, 0, 0]);
     this.strides = attributes.getInts('strides', [1, 1]);
     this.activation = attributes.getString('__internal_activation', '');
+    this.clipMax = attributes.getFloat('__clip_max', 3.402823e+38);
+    this.clipMin = attributes.getFloat('__clip_min', -3.402823e+38);
   }
 
   checkInputs(inputs: Tensor[]): boolean {
@@ -90,4 +92,6 @@ export abstract class Conv implements Operator {
   protected pads: number[];
   protected strides: number[];
   protected activation: string;
+  protected clipMin: number;
+  protected clipMax: number;
 }

js/web/lib/onnxjs/ops/matmul.ts

@@ -11,6 +11,8 @@ export abstract class MatMul implements Operator {
 
   initialize(attributes: Attribute): void {
     this.activation = attributes.getString('__internal_activation', '');
+    this.clipMax = attributes.getFloat('__clip_max', 3.402823e+38);
+    this.clipMin = attributes.getFloat('__clip_min', -3.402823e+38);
   }
 
   checkInputs(inputs: Tensor[]): boolean {
@@ -41,4 +43,6 @@ export abstract class MatMul implements Operator {
     return true;
   }
   protected activation: string;
+  protected clipMin: number;
+  protected clipMax: number;
 }