From fa722d208b4af6f1c335273a1e465ae0f9d892dc Mon Sep 17 00:00:00 2001 From: Du Li Date: Wed, 4 Aug 2021 20:46:08 -0700 Subject: [PATCH] [js/web] adding webgl pointwise conv kernel (#8418) --- js/web/lib/onnxjs/backends/webgl/ops/conv.ts | 69 ++++++++++++++- .../onnxjs/backends/webgl/ops/matmul-pack.ts | 33 +------ .../lib/onnxjs/backends/webgl/ops/matmul.ts | 64 +++++++++++++- .../lib/onnxjs/backends/webgl/ops/reshape.ts | 11 ++- js/web/test/data/ops/conv.jsonc | 87 +++++++++++++++++++ 5 files changed, 226 insertions(+), 38 deletions(-) diff --git a/js/web/lib/onnxjs/backends/webgl/ops/conv.ts b/js/web/lib/onnxjs/backends/webgl/ops/conv.ts index a76298a703..a41f590672 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/conv.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/conv.ts @@ -5,7 +5,7 @@ import {Attribute} from '../../../attribute'; import {Logger} from '../../../instrument'; import {Conv} from '../../../ops/conv'; import {Tensor} from '../../../tensor'; -import {PoolConvUtil} from '../../../util'; +import {assert, PoolConvUtil} from '../../../util'; import {getGlsl} from '../glsl-source'; import {WebGLInferenceHandler} from '../inference-handler'; import {Artifact, ProgramInfo, RunData, TextureLayout, WebGLOperator} from '../types'; @@ -13,17 +13,21 @@ import {WebGLContext} from '../webgl-context'; import {WebGLConvPacked} from './conv-pack'; import {getActicationSnippet} from './fuse-utils'; +import {WebGLUnpackedMatMul} from './matmul'; +import {reshape} from './reshape'; export class WebGLConv extends Conv { unpackedGroupedConvImpl: WebGLUnpackedGroupedConv; unpackedConvImpl: WebGLUnpackedConv; packedConvImpl: WebGLConvPacked; + unpackedPointwiseConvImpl: WebGLUnpackedPointwiseConv; constructor() { super(); this.unpackedGroupedConvImpl = new WebGLUnpackedGroupedConv(); this.unpackedConvImpl = new WebGLUnpackedConv(); this.packedConvImpl = new WebGLConvPacked(); + this.unpackedPointwiseConvImpl = new WebGLUnpackedPointwiseConv(); } initialize(attributes: Attribute): void { @@ -31,6 +35,7 @@ export class WebGLConv extends Conv { this.unpackedGroupedConvImpl.initialize(attributes); this.unpackedConvImpl.initialize(attributes); this.packedConvImpl.initialize(attributes); + this.unpackedPointwiseConvImpl.initialize(attributes); } run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] { @@ -38,7 +43,9 @@ export class WebGLConv extends Conv { 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 && !isPointwise) { + } else if (isPointwise && packMode) { + return this.unpackedPointwiseConvImpl.run(inferenceHandler, inputs); + } else if (packMode && inputs[0].dims.length === 4 && inputs[0].dims[0] === 1) { return this.packedConvImpl.run(inferenceHandler, inputs); } else { return this.unpackedConvImpl.run(inferenceHandler, inputs); @@ -62,6 +69,64 @@ export class WebGLConv extends Conv { } } +export class WebGLUnpackedPointwiseConv extends Conv { + private programInfo: ProgramInfo[]; + private artifacts: Artifact[]; + private matmul = new WebGLUnpackedMatMul(); + private outputShape: number[]; + + run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] { + const programManager = inferenceHandler.session.programManager; + const xShape = inputs[0].dims.slice(); + const wShape = inputs[1].dims.slice(); + PoolConvUtil.adjustPadsBasedOnAutoPad( + inputs[0].dims, this.strides, this.dilations, this.kernelShape, this.pads, this.autoPad); + Logger.verbose( + 'Conv', + `autpPad:${this.autoPad}, dilations:${this.dilations}, group:${this.group}, kernelShape:${ + this.kernelShape}, pads:${this.pads}, strides:${this.strides}`); + + if (!this.outputShape) { + this.outputShape = WebGLConv.calcOutputShape(xShape, wShape, this.dilations, this.pads, this.strides); + } + + 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); + } + const reshapedX = reshape(inferenceHandler, inputs[0], [xShape[1], xShape[2] * xShape[3]]); + const reshapedW = reshape(inferenceHandler, inputs[1], [wShape[0], wShape[1]]); + const hasBias = (inputs.length === 3); + + if (!this.artifacts) { + this.artifacts = []; + this.programInfo = []; + this.programInfo[0] = hasBias ? + this.matmul.createProgramInfo(inferenceHandler, [reshapedW, reshapedX, inputs[2]]) : + this.matmul.createProgramInfo(inferenceHandler, [reshapedW, reshapedX]); + this.artifacts[0] = programManager.build(this.programInfo[0]); + } + + assert(this.artifacts.length === 1, () => 'expect matmul artifacts was created'); + const runDataMatmul = this.matmul.createRunData( + inferenceHandler, this.programInfo[0], hasBias ? [reshapedW, reshapedX, inputs[2]] : [reshapedW, reshapedX]); + programManager.run(this.artifacts[0], runDataMatmul); + const matmulOutput = runDataMatmul.outputTextureData.tensor; + + const res = [reshape(inferenceHandler, matmulOutput, this.outputShape)]; + // eslint-disable-next-line no-console + // console.log( + // 'res[0]=', res[0].data[0], ' res[1]=', res[0].data[1], ' res[1]=', res[0].data[2], ' res[3]=', + // res[0].data[3]); + return res; + } +} + export class WebGLUnpackedGroupedConv extends Conv implements WebGLOperator { run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] { return inferenceHandler.run(this, inputs); 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 f2dfae11fc..a3e540d1fd 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts @@ -11,6 +11,7 @@ import {ProgramInfo, RunData, WebGLOperator} from '../types'; import {getCoordsDataType} from '../utils'; import {getActicationSnippet} from './fuse-utils'; +import {getBiasForMatmul} from './matmul'; export class WebGLMatMulPacked extends MatMul implements WebGLOperator { run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] { @@ -98,38 +99,6 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator { } } -function getBiasForMatmul( - coordsDataType: string, allGlChannels: readonly string[], inShape: readonly number[], - outShape: readonly number[]): string { - let unpackedCoordsSnippet = ''; - const inRank = inShape.length; - const outRank = outShape.length; - const rankDiff = outRank - inRank; - if (outRank < 2 && inRank > 0) { - unpackedCoordsSnippet = 'coords'; - } else { - unpackedCoordsSnippet = inShape.map((s, i) => `coords.${allGlChannels[i + rankDiff]}`).join(', '); - } - const broadcastDims = BroadcastUtil.getBroadcastDims(inShape, outShape); - const coordsSnippet = broadcastDims.map(d => `coords.${allGlChannels[d + rankDiff]} = 0;`).join('\n'); - const inSize = ShapeUtil.size(inShape); - const isInputScalar = inSize === 1; - let output = 'vec4(outputValue.xx, outputValue.yy)'; - if (isInputScalar) { - output = 'vec4(outputValue.x)'; - } - const getBiasForMatmulSource = ` - vec4 getBiasForMatmul() { - ${coordsDataType} coords = getOutputCoords(); - ${coordsSnippet} - vec4 outputValue = getBias(${unpackedCoordsSnippet}); - return ${output}; - - }`; - - return getBiasForMatmulSource; -} - function getBcastSamplerForMatmul( coordsDataType: string, allGlChannels: readonly string[], inputs: Tensor[], outShape: readonly number[]): string { let unpackedACoordsSnippet = []; diff --git a/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts b/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts index 89aed2444c..45dc9fa50f 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts @@ -3,9 +3,11 @@ import {MatMul} from '../../../ops/matmul'; import {Tensor} from '../../../tensor'; -import {BroadcastUtil} from '../../../util'; +import {BroadcastUtil, ShapeUtil} from '../../../util'; import {WebGLInferenceHandler} from '../inference-handler'; import {ProgramInfo, RunData, WebGLOperator} from '../types'; +import {getCoordsDataType} from '../utils'; +import {getActicationSnippet} from './fuse-utils'; import {WebGLMatMulPacked} from './matmul-pack'; export class WebGLMatMul extends MatMul implements WebGLOperator { @@ -57,17 +59,34 @@ export class WebGLUnpackedMatMul extends MatMul implements WebGLOperator { return inferenceHandler.run(this, inputs); } createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo { + const hasBias = inputs.length > 2; + const processBias = hasBias ? 'value += getBiasForMatmul();' : ''; const aShape = inputs[0].dims; const bShape = inputs[1].dims; const outputShape = BroadcastUtil.calcShape(aShape, bShape, true); if (!outputShape) { throw new Error('Can\'t use matmul on the given tensors'); } + const coordsDataType = getCoordsDataType(outputShape.length); + const allGlChannels = ['x', 'y', 'z', 'w', 'u', 'v']; + + const {activationFunction, applyActivation} = getActicationSnippet(this.activation); + const additionalVars = this.activation === 'Clip' ? ` + float min = float(${this.clipMin}); + float max = float(${this.clipMax});` : + ''; + const getBiasForMatmulSnippet = + hasBias ? `${getBiasForMatmul(coordsDataType, allGlChannels, inputs[2].dims, outputShape, false)}` : ''; + const rank = outputShape.length; const arank = aShape.length; const brank = bShape.length; const sharedDim = aShape[aShape.length - 1]; const shaderSource = ` + ${additionalVars} + ${activationFunction} + ${getBiasForMatmulSnippet} + float process(int indices[${rank}]) { int a[${arank}]; int b[${brank}]; @@ -80,12 +99,15 @@ export class WebGLUnpackedMatMul extends MatMul implements WebGLOperator { b[${brank - 2}] = k; value += _A(a) * _B(b); } + ${processBias} + ${applyActivation} + return value; }`; return { inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t)), outputLayout: handler.createTextureLayoutFromShape(outputShape), - samplers: ['A', 'B'], + samplers: hasBias ? ['A', 'B', 'Bias'] : ['A', 'B'], shaderSource, }; } @@ -98,3 +120,41 @@ export class WebGLUnpackedMatMul extends MatMul implements WebGLOperator { }; } } + +export function getBiasForMatmul( + coordsDataType: string, allGlChannels: readonly string[], inShape: readonly number[], outShape: readonly number[], + isPacked = true): string { + let unpackedCoordsSnippet = ''; + const inRank = inShape.length; + const outRank = outShape.length; + const rankDiff = outRank - inRank; + if (outRank < 2 && inRank > 0) { + unpackedCoordsSnippet = 'coords'; + } else { + unpackedCoordsSnippet = inShape.map((s, i) => `coords.${allGlChannels[i + rankDiff]}`).join(', '); + } + const broadcastDims = BroadcastUtil.getBroadcastDims(inShape, outShape); + const coordsSnippet = broadcastDims.map(d => `coords.${allGlChannels[d + rankDiff]} = 0;`).join('\n'); + const inSize = ShapeUtil.size(inShape); + const isInputScalar = inSize === 1; + let output = 'vec4(outputValue.xx, outputValue.yy)'; + if (isInputScalar) { + output = 'vec4(outputValue.x)'; + } + const getBiasForMatmulSource = isPacked ? ` +vec4 getBiasForMatmul() { + ${coordsDataType} coords = getOutputCoords(); + ${coordsSnippet} + vec4 outputValue = getBias(${unpackedCoordsSnippet}); + return ${output}; +}` : + ` +float getBiasForMatmul() { + ${coordsDataType} coords = getOutputCoords(); + ${coordsSnippet} + return getBias(coords.x); +} +`; + + return getBiasForMatmulSource; +} diff --git a/js/web/lib/onnxjs/backends/webgl/ops/reshape.ts b/js/web/lib/onnxjs/backends/webgl/ops/reshape.ts index 6ae4117d1d..1825a824c2 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/reshape.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/reshape.ts @@ -7,8 +7,10 @@ import {ShapeUtil} from '../../../util'; import {WebGLInferenceHandler} from '../inference-handler'; import {TextureLayout} from '../types'; import {getPackedShape} from '../utils'; + import {WebGLReshapePacked} from './reshape-packed'; + export class WebGLReshape extends Reshape { packedImpl: WebGLReshapePacked; constructor() { @@ -28,10 +30,15 @@ export class WebGLReshape extends Reshape { export function reshape( inferenceHandler: WebGLInferenceHandler, input: Tensor, reshapedDims: readonly number[]): Tensor { - const inputTD = inferenceHandler.getOrCreateTextureData(input); + let inputTD = inferenceHandler.getOrCreateTextureData(input); let packedShape = reshapedDims; + if (inputTD.channels === 4) { - packedShape = getPackedShape(reshapedDims); + if (!inputTD.isPacked) { + packedShape = getPackedShape(reshapedDims); + } else { + inputTD = inferenceHandler.unpack(inputTD); + } } const newTextureLayout: TextureLayout = { channels: inputTD.channels, diff --git a/js/web/test/data/ops/conv.jsonc b/js/web/test/data/ops/conv.jsonc index 1950606c82..35f7360454 100644 --- a/js/web/test/data/ops/conv.jsonc +++ b/js/web/test/data/ops/conv.jsonc @@ -416,5 +416,92 @@ ] } ] + }, + + { + "name": "conv - pointwise", + "operator": "Conv", + "attributes": [{ "name": "kernel_shape", "data": [1, 1], "type": "ints" }], + "cases": [ + { + "name": "T[0]", + "inputs": [ + { + "data": [ + 0.0, + 1.0, + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 7.0, + 8.0, + 9.0, + 10.0, + 11.0, + 12.0, + 13.0, + 14.0, + 15.0, + 16.0, + 17.0, + 18.0, + 19.0, + 20.0, + 21.0, + 22.0, + 23.0, + 24.0, + 25.0, + 26.0, + 27.0, + 28.0, + 29.0, + 30.0, + 31.0 + ], + "dims": [1, 8, 2, 2], + "type": "float32" + }, + { + "data":[ + 0.0, + 1.0, + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 7.0, + 8.0, + 9.0, + 10.0, + 11.0, + 12.0, + 13.0, + 14.0, + 15.0 + ], + "dims":[2, 8, 1, 1], + "type":"float32" + }, + { + "data": [0.5, 0.4], + "dims": [2], + "type": "float32" + } + ], + "outputs":[ + { + "data": [ + 560.5,588.5,616.5,644.5,1456.4,1548.4,1640.4,1732.4 + ], + "dims": [1, 2, 2, 2], + "type": "float32" + } + ] + } + ] } ]