From db88f3059c6a16e35326992fc0f1f73b0a020771 Mon Sep 17 00:00:00 2001 From: Tixxx Date: Wed, 23 Jun 2021 09:55:19 -0700 Subject: [PATCH] [js] fixing broadcast issues in pack mode (#8090) * fixing broadcast issues in pack mode * improved bcast logic for matmul * removed TODO * rebased from master --- .../backends/webgl/glsl-coordinate-lib.ts | 9 +- .../backends/webgl/glsl-shape-utils-lib.ts | 2 +- .../onnxjs/backends/webgl/ops/binary-op.ts | 35 ++++-- .../onnxjs/backends/webgl/ops/matmul-pack.ts | 113 ++++++++++++++++-- .../lib/onnxjs/backends/webgl/ops/matmul.ts | 8 +- .../backends/webgl/test-matmul-packed.ts | 59 +++++++-- 6 files changed, 189 insertions(+), 37 deletions(-) diff --git a/js/web/lib/onnxjs/backends/webgl/glsl-coordinate-lib.ts b/js/web/lib/onnxjs/backends/webgl/glsl-coordinate-lib.ts index 097d3a46c4..0122ad16a1 100644 --- a/js/web/lib/onnxjs/backends/webgl/glsl-coordinate-lib.ts +++ b/js/web/lib/onnxjs/backends/webgl/glsl-coordinate-lib.ts @@ -654,7 +654,7 @@ export class CoordsGlslLib extends GlslLib { if (inRank === 1 && !isInputScalar && !isOutputScalar) { output = ` - return vec4(outputValue.xx, outputValue.yy); + return vec4(outputValue.xy, outputValue.xy); `; } else if (isInputScalar && !isOutputScalar) { if (outRank === 1) { @@ -679,9 +679,16 @@ export class CoordsGlslLib extends GlslLib { output = 'return vec4(outputValue.xx, outputValue.zz);'; } } + + const swapLastDimsSnippet = ` + int lastDim = coords.${fields[outRank - 1]}; + coords.${fields[outRank - 1]} = coords.${fields[outRank - 2]}; + coords.${fields[outRank - 2]} = lastDim; + `; const source = ` vec4 ${funcName}() { ${type} coords = getOutputCoords(); + ${swapLastDimsSnippet} ${coordsSnippet} vec4 outputValue = ${texFuncSnippet}(${unpackedCoordsSnippet}); ${output} diff --git a/js/web/lib/onnxjs/backends/webgl/glsl-shape-utils-lib.ts b/js/web/lib/onnxjs/backends/webgl/glsl-shape-utils-lib.ts index f337dace84..8045ec9f84 100644 --- a/js/web/lib/onnxjs/backends/webgl/glsl-shape-utils-lib.ts +++ b/js/web/lib/onnxjs/backends/webgl/glsl-shape-utils-lib.ts @@ -28,7 +28,7 @@ export class ShapeUtilsGlslLib extends GlslLib { const outputRank = programInfo.outputLayout.shape.length; const result: {[name: string]: GlslLibRoutine} = {}; this.context.programInfo.samplers.forEach((name, i) => { - const shape = programInfo.inputLayouts[i].shape; + const shape = programInfo.inputLayouts[i].unpackedShape; if (shape.length <= outputRank) { const rank = shape.length; const dimOffset = outputRank - rank; diff --git a/js/web/lib/onnxjs/backends/webgl/ops/binary-op.ts b/js/web/lib/onnxjs/backends/webgl/ops/binary-op.ts index 509f25a6a7..5264c55a5a 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/binary-op.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/binary-op.ts @@ -22,32 +22,41 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator { } createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo { const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims); + const outputShape = isBroadcast ? BroadcastUtil.calcShape(inputs[0].dims, inputs[1].dims, false) : inputs[0].dims; + if (!outputShape) { + throw new Error('Can\'t perform binary op on the given tensors'); + } - // TODO fix bcast in packed mode. if (this.usePackedTexture === undefined) { - this.usePackedTexture = !isBroadcast && handler.session.pack; + this.usePackedTexture = handler.session.pack; } const inputLayouts = this.usePackedTexture ? inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)) : inputs.map(t => handler.getOrCreateTextureLayout(t)); - const ouputLayout = this.usePackedTexture ? - handler.createTextureLayoutFromShape(inputs[0].dims, 4, inputs[0].dims, {isPacked: true, reverseWH: true}) : + let outputLayout = this.usePackedTexture ? + handler.createTextureLayoutFromShape(inputs[0].dims, 4, outputShape, {isPacked: true, reverseWH: true}) : handler.createTextureLayoutFromShape(inputs[0].dims); + const glsl = getGlsl(handler.session.backend.glContext.version); if (isBroadcast) { - const outputShape = BroadcastUtil.calcShape(inputs[0].dims, inputs[1].dims, false); - if (!outputShape) { - throw new Error('Can\'t perform binary op on the given tensors'); - } const outputRank = outputShape.length; const aRank = inputs[0].dims.length !== 0 ? inputs[0].dims.length : 1; 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 = ` + const shaderSource = this.usePackedTexture ? ` + ${this.glslFunc.body} + void main() { + vec4 a = getAAtOutCoords(); + vec4 b = getBAtOutCoords(); + + vec4 result = ${this.glslFunc.name}(a, b); + + ${glsl.output} = result; + }` : + ` ${this.glslFunc.body} float process(int indices[${outputRank}]) { int aindices[${aRank}]; @@ -56,11 +65,12 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator { ${bBcast} return ${this.glslFunc.name}(_A(aindices), _B(bindices)); }`; - const outputLayout = this.usePackedTexture ? + outputLayout = this.usePackedTexture ? handler.createTextureLayoutFromShape(outputShape, 4, outputShape, {isPacked: true, reverseWH: true}) : handler.createTextureLayoutFromShape(outputShape); return { + hasMain: this.usePackedTexture, inputLayouts, outputLayout, samplers: ['A', 'B'], @@ -69,7 +79,6 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator { expectPackedOutputs: this.usePackedTexture }; } - const glsl = getGlsl(handler.session.backend.glContext.version); const shaderSource = ` ${this.glslFunc.body} void main() { @@ -84,7 +93,7 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator { return { hasMain: true, inputLayouts, - outputLayout: ouputLayout, + outputLayout, samplers: ['A', 'B'], shaderSource, expectPackedInputs: true, 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 ecb80d39c4..f2dfae11fc 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts @@ -4,6 +4,7 @@ import {MatMul} from '../../../ops/matmul'; import {Tensor} from '../../../tensor'; import {BroadcastUtil} from '../../../util'; +import {ShapeUtil} from '../../../util'; import {getGlsl} from '../glsl-source'; import {WebGLInferenceHandler} from '../inference-handler'; import {ProgramInfo, RunData, WebGLOperator} from '../types'; @@ -17,10 +18,11 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator { } createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo { const hasBias = inputs.length > 2; - const processBias = hasBias ? 'value += getBiasAtOutCoords();' : ''; + const processBias = hasBias ? 'value += getBiasForMatmul();' : ''; const aShape = inputs[0].dims; const bShape = inputs[1].dims; const outputShape = BroadcastUtil.calcShape(aShape, bShape, true); + const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims); if (!outputShape) { throw new Error('Can\'t use matmul on the given tensors'); @@ -39,19 +41,33 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator { float max = float(${this.clipMax});` : ''; const {activationFunction, applyActivation} = getActicationSnippet(this.activation); + + const getBiasForMatmulSnippet = + hasBias ? `${getBiasForMatmul(coordsDataType, allGlChannels, inputs[2].dims, outputShape)}` : ''; + + const getBcastedSamplerForMatmulSnippet = + isBroadcast ? `${getBcastSamplerForMatmul(coordsDataType, allGlChannels, inputs, outputShape)}` : ''; + + const getSamplerAInLoopSnippet = isBroadcast ? 'getAAtOutCoordsMatmul(i)' : `getA(${getA(allGlChannels, aRank)})`; + const getSamplerBInLoopSnippet = isBroadcast ? 'getBAtOutCoordsMatmul(i)' : `getB(${getB(allGlChannels, bRank)})`; + const getOutputCoordsSnippet = isBroadcast ? '' : `${coordsDataType} rc = getOutputCoords(); + int lastDim = rc.${allGlChannels[outRank - 1]}; + rc.${allGlChannels[outRank - 1]} = rc.${allGlChannels[outRank - 2]}; + rc.${allGlChannels[outRank - 2]} = lastDim; +`; const shaderSource = ` ${additionalVars} + ${getBcastedSamplerForMatmulSnippet} + ${getBiasForMatmulSnippet} ${activationFunction} void main() { - ${coordsDataType} rc = getOutputCoords(); - int lastDim = rc.${allGlChannels[outRank - 1]}; - rc.${allGlChannels[outRank - 1]} = rc.${allGlChannels[outRank - 2]}; - rc.${allGlChannels[outRank - 2]} = lastDim; + ${getOutputCoordsSnippet} vec4 value = vec4(0); for (int i = 0; i < ${sharedDimIndex}; i++) { - vec4 a = getA(${getA(allGlChannels, aRank)}); - vec4 b = getB(${getB(allGlChannels, bRank)}); + vec4 a = ${getSamplerAInLoopSnippet}; + vec4 b = ${getSamplerBInLoopSnippet}; + value += (a.rrbb * b.rgrg); value += (a.ggaa * b.baba); } @@ -82,6 +98,89 @@ 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 = []; + let unpackedBCoordsSnippet = []; + + const inAShape = inputs[0].dims; + const inBShape = inputs[1].dims; + + const inARank = inAShape.length; + const inBRank = inBShape.length; + + const outRank = outShape.length; + const rankADiff = outRank - inARank; + const rankBDiff = outRank - inBRank; + + unpackedACoordsSnippet = inAShape.map((s, i) => `coords.${allGlChannels[i + rankADiff]}`); + unpackedACoordsSnippet[inARank - 1] = 'i*2'; + unpackedACoordsSnippet.join(', '); + unpackedBCoordsSnippet = inBShape.map((s, i) => `coords.${allGlChannels[i + rankBDiff]}`); + unpackedBCoordsSnippet[inBRank - 2] = 'i*2'; + unpackedBCoordsSnippet.join(', '); + + const broadcastADims = BroadcastUtil.getBroadcastDims(inAShape, outShape); + const broadcastBDims = BroadcastUtil.getBroadcastDims(inBShape, outShape); + + const coordsASnippet = broadcastADims.map(d => `coords.${allGlChannels[d + rankADiff]} = 0;`).join('\n'); + const coordsBSnippet = broadcastBDims.map(d => `coords.${allGlChannels[d + rankBDiff]} = 0;`).join('\n'); + const swapDimSnippet = `int lastDim = coords.${allGlChannels[outRank - 1]}; + coords.${allGlChannels[outRank - 1]} = coords.${allGlChannels[outRank - 2]}; + coords.${allGlChannels[outRank - 2]} = lastDim;`; + + const getBcastSamplerMatmulSource = ` +vec4 getAAtOutCoordsMatmul(int i) { + ${coordsDataType} coords = getOutputCoords(); + ${swapDimSnippet} + ${coordsASnippet} + vec4 outputValue = getA(${unpackedACoordsSnippet}); + return outputValue; +} + +vec4 getBAtOutCoordsMatmul(int i) { + ${coordsDataType} coords = getOutputCoords(); + ${swapDimSnippet} + ${coordsBSnippet} + vec4 outputValue = getB(${unpackedBCoordsSnippet}); + return outputValue; +}`; + + return getBcastSamplerMatmulSource; +} + function getA(allGlChannels: string[], rank: number): string { let res = ''; for (let i = 0; i < rank - 2; i++) { diff --git a/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts b/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts index 4188352ddf..89aed2444c 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/matmul.ts @@ -3,7 +3,7 @@ import {MatMul} from '../../../ops/matmul'; import {Tensor} from '../../../tensor'; -import {BroadcastUtil, ShapeUtil} from '../../../util'; +import {BroadcastUtil} from '../../../util'; import {WebGLInferenceHandler} from '../inference-handler'; import {ProgramInfo, RunData, WebGLOperator} from '../types'; import {WebGLMatMulPacked} from './matmul-pack'; @@ -21,8 +21,7 @@ export class WebGLMatMul extends MatMul implements WebGLOperator { run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] { if (this.usePackedTexture === undefined) { - const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims); - this.usePackedTexture = !isBroadcast && inferenceHandler.session.pack; + this.usePackedTexture = inferenceHandler.session.pack; } if (this.usePackedTexture) { @@ -34,8 +33,7 @@ export class WebGLMatMul extends MatMul implements WebGLOperator { createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo { if (this.usePackedTexture === undefined) { - const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims); - this.usePackedTexture = !isBroadcast && handler.session.pack; + this.usePackedTexture = handler.session.pack; } if (this.usePackedTexture && inputs[0].dims.length > 1) { diff --git a/js/web/test/unittests/backends/webgl/test-matmul-packed.ts b/js/web/test/unittests/backends/webgl/test-matmul-packed.ts index f9757c7d67..256e7779fc 100644 --- a/js/web/test/unittests/backends/webgl/test-matmul-packed.ts +++ b/js/web/test/unittests/backends/webgl/test-matmul-packed.ts @@ -7,6 +7,7 @@ import {env} from 'onnxruntime-common'; import {Backend, InferenceHandler, resolveBackend, SessionHandler} from '../../../../lib/onnxjs/backend'; import {WebGLInferenceHandler} from '../../../../lib/onnxjs/backends/webgl/inference-handler'; import {WebGLMatMulPacked} from '../../../../lib/onnxjs/backends/webgl/ops/matmul-pack'; +import {TextureData} from '../../../../lib/onnxjs/backends/webgl/types'; import {Profiler} from '../../../../lib/onnxjs/instrument'; import {Tensor} from '../../../../lib/onnxjs/tensor'; import {ShapeUtil} from '../../../../lib/onnxjs/util'; @@ -129,6 +130,19 @@ function getTestData(): TestData[] { rawInputA: new Float32Array([1, 2, 4, 5, 3, 0, 6, 0, 1, 2, 4, 5, 3, 0, 6, 0]), rawInputB: new Float32Array([1, 2, 3, 4, 5, 6, 0, 0]), }, + { + elementCountA: 12, + elementCountB: 6, + inputShapeA: [1, 2, 2, 3], + inputShapeB: [1, 1, 1, 3, 2], + outputShape: [1, 1, 2, 2, 2], + inputTextureShapeA: [2, 2], + inputTextureShapeB: [1, 2], + outputTextureShape: [2, 1], + expectedOutput: new Float32Array([22, 28, 49, 64, 22, 28, 49, 64]), + rawInputA: new Float32Array([1, 2, 4, 5, 3, 0, 6, 0, 1, 2, 4, 5, 3, 0, 6, 0]), + rawInputB: new Float32Array([1, 2, 3, 4, 5, 6, 0, 0]), + }, ]; } @@ -173,6 +187,9 @@ describe('#UnitTest# - packed matmul - Tensor matmul', () => { const inputDataB = createAscendingArray(elementCountB); const inputTensorA = new Tensor(inputTensorShapeA, 'float32', undefined, undefined, inputDataA); const inputTensorB = new Tensor(inputTensorShapeB, 'float32', undefined, undefined, inputDataB); + const biasTensor = testData.biasValue ? + new Tensor([1], 'float32', undefined, undefined, new Float32Array([testData.biasValue])) : + undefined; // manually creat packed texture from inputTensor, and insert in cache const gl = webglInferenceHandler.session.textureManager.glContext.gl; @@ -185,6 +202,12 @@ describe('#UnitTest# - packed matmul - Tensor matmul', () => { webglInferenceHandler.session.textureManager.glContext, testData.rawInputB ? testData.rawInputB : inputDataB, gl.RGBA, inputTextureShapeB[0], inputTextureShapeB[1]); + const webglTextureBias = biasTensor && testData.biasValue ? + createTextureFromArray( + webglInferenceHandler.session.textureManager.glContext, new Float32Array([testData.biasValue, 0, 0, 0]), + gl.RGBA, 1, 1) : + undefined; + webglInferenceHandler.session.textureManager.glContext.checkError(); const packedShapeA = inputTextureShapeA; const textureDataA = { @@ -212,15 +235,25 @@ describe('#UnitTest# - packed matmul - Tensor matmul', () => { texture: webglTextureB! }; + const packedShapeBias = [1]; webglInferenceHandler.setTextureData(inputTensorA.dataId, textureDataA, true); webglInferenceHandler.setTextureData(inputTensorB.dataId, textureDataB, true); + if (biasTensor && webglTextureBias) { + const textureDataBias: TextureData = { + width: 1, + height: 1, + channels: 4 as const, + isPacked: true, + shape: packedShapeBias, + strides: ShapeUtil.computeStrides(packedShapeBias), + unpackedShape: [1], + tensor: biasTensor, + texture: webglTextureBias + }; - const inputList = testData.biasValue ? - [ - inputTensorA, inputTensorB, - new Tensor([1], 'float32', undefined, undefined, new Float32Array([testData.biasValue])) - ] : - [inputTensorA, inputTensorB]; + webglInferenceHandler.setTextureData(biasTensor.dataId, textureDataBias, true); + } + const inputList = biasTensor ? [inputTensorA, inputTensorB, biasTensor] : [inputTensorA, inputTensorB]; // compile shader code const programInfo = op.createProgramInfo(inferenceHandler! as WebGLInferenceHandler, inputList); @@ -237,14 +270,20 @@ describe('#UnitTest# - packed matmul - Tensor matmul', () => { // verify result. const expectedOutput = testData.expectedOutput; expect(result).to.not.equal(null); - let batchMultiplier = 1; + let batchMultiplierA = 1; + let batchMultiplierB = 1; + if (testData.inputShapeA.length > 2) { - batchMultiplier = testData.inputShapeA[0]; + for (let i = 0; i < testData.inputShapeA.length - 2; i++) { + batchMultiplierA *= testData.inputShapeA[i]; + } } if (testData.inputShapeB.length > 2) { - batchMultiplier = Math.max(batchMultiplier, testData.inputShapeB[0]); + for (let i = 0; i < testData.inputShapeB.length - 2; i++) { + batchMultiplierB *= testData.inputShapeB[i]; + } } - + const batchMultiplier = Math.max(batchMultiplierA, batchMultiplierB); expect(result).to.have.lengthOf( batchMultiplier * testData.inputShapeA[testData.inputShapeA.length - 2] * testData.inputShapeB[testData.inputShapeB.length - 1]);