From 352d560fd586c35f96495092e75234b644ec1d8e Mon Sep 17 00:00:00 2001 From: Du Li Date: Mon, 21 Jun 2021 16:30:12 -0700 Subject: [PATCH] Adding Conv+Clip fusion (#8102) --- js/web/lib/onnxjs/backends/webgl/ops/clip.ts | 16 ++++++++++++++++ .../lib/onnxjs/backends/webgl/ops/conv-pack.ts | 4 ++++ js/web/lib/onnxjs/backends/webgl/ops/conv.ts | 16 ++++++++++++---- .../lib/onnxjs/backends/webgl/ops/fuse-utils.ts | 12 ++++++++++-- .../lib/onnxjs/backends/webgl/ops/matmul-pack.ts | 16 ++++++++++------ js/web/lib/onnxjs/graph.ts | 14 +++++++++++--- js/web/lib/onnxjs/ops/conv.ts | 4 ++++ js/web/lib/onnxjs/ops/matmul.ts | 4 ++++ 8 files changed, 71 insertions(+), 15 deletions(-) diff --git a/js/web/lib/onnxjs/backends/webgl/ops/clip.ts b/js/web/lib/onnxjs/backends/webgl/ops/clip.ts index 13bbed7467..a4e5cc4964 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/clip.ts +++ b/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}; +} diff --git a/js/web/lib/onnxjs/backends/webgl/ops/conv-pack.ts b/js/web/lib/onnxjs/backends/webgl/ops/conv-pack.ts index 68c962b5e5..fc8105c73c 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/conv-pack.ts +++ b/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 diff --git a/js/web/lib/onnxjs/backends/webgl/ops/conv.ts b/js/web/lib/onnxjs/backends/webgl/ops/conv.ts index 8fdb052ceb..a76298a703 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/conv.ts +++ b/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]; diff --git a/js/web/lib/onnxjs/backends/webgl/ops/fuse-utils.ts b/js/web/lib/onnxjs/backends/webgl/ops/fuse-utils.ts index 115cae72bc..7a8808350f 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/fuse-utils.ts +++ b/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}; } diff --git a/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts b/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts index fc71e5991a..ecb80d39c4 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts +++ b/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', diff --git a/js/web/lib/onnxjs/graph.ts b/js/web/lib/onnxjs/graph.ts index af281d1047..b5bd5e39c5 100644 --- a/js/web/lib/onnxjs/graph.ts +++ b/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]); } } } } -} \ No newline at end of file +} diff --git a/js/web/lib/onnxjs/ops/conv.ts b/js/web/lib/onnxjs/ops/conv.ts index b146ad69ed..85b9e49cc4 100644 --- a/js/web/lib/onnxjs/ops/conv.ts +++ b/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; } diff --git a/js/web/lib/onnxjs/ops/matmul.ts b/js/web/lib/onnxjs/ops/matmul.ts index 0ad0aeb094..e34880e625 100644 --- a/js/web/lib/onnxjs/ops/matmul.ts +++ b/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; }