mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
[js] fixing broadcast issues in pack mode (#8090)
* fixing broadcast issues in pack mode * improved bcast logic for matmul * removed TODO * rebased from master
This commit is contained in:
Tixxx
GitHub
parent
cbdd59dae9
commit
db88f3059c
6 changed files with 189 additions and 37 deletions
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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++) {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
|
|
|
|||
Loading…
Reference in a new issue