mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-06-03 23:49:44 +00:00
312 lines
13 KiB
TypeScript
312 lines
13 KiB
TypeScript
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
import {InferenceHandler} from '../../backend';
|
|
import {Logger} from '../../instrument';
|
|
import {Tensor} from '../../tensor';
|
|
import {ShapeUtil} from '../../util';
|
|
import {WebGLPack} from './ops/pack';
|
|
|
|
import {WebGLUint8Encode} from './ops/uint8-encode';
|
|
import {WebGLUnpack} from './ops/unpack';
|
|
import {WebGLSessionHandler} from './session-handler';
|
|
import {Encoder} from './texture-data-encoder';
|
|
import {WidthHeightPrefs} from './texture-layout-strategy';
|
|
import {Artifact, RunData, TextureData, TextureLayout, WebGLOperator} from './types';
|
|
import {getPackedShape} from './utils';
|
|
|
|
export class WebGLInferenceHandler implements InferenceHandler {
|
|
private packedTextureDataCache: Map<Tensor.Id, TextureData>;
|
|
private unpackedTextureDataCache: Map<Tensor.Id, TextureData>;
|
|
private pack2unpackMap: Map<Tensor.Id, Tensor.Id>;
|
|
private unpack2packMap: Map<Tensor.Id, Tensor.Id>;
|
|
constructor(public session: WebGLSessionHandler) {
|
|
this.packedTextureDataCache = new Map();
|
|
this.unpackedTextureDataCache = new Map();
|
|
|
|
this.pack2unpackMap = new Map();
|
|
this.unpack2packMap = new Map();
|
|
}
|
|
|
|
run(op: WebGLOperator, inputs: Tensor[]): Tensor[] {
|
|
let artifact = this.session.programManager.getArtifact(op);
|
|
if (!artifact) {
|
|
const programInfo = op.createProgramInfo(this, inputs);
|
|
if (!programInfo.name) {
|
|
programInfo.name = op.constructor?.name;
|
|
}
|
|
artifact = this.session.programManager.build(programInfo);
|
|
this.session.programManager.setArtifact(op, artifact);
|
|
}
|
|
const runData = op.createRunData(this, artifact.programInfo, inputs);
|
|
this.runProgram(artifact, runData);
|
|
return [runData.outputTextureData.tensor];
|
|
}
|
|
|
|
checkAndUpdateTextureForm(artifact: Artifact, runData: RunData) {
|
|
// pack/unpack inputs
|
|
for (let i = 0; i < runData.inputTextureDatas.length; ++i) {
|
|
const input = runData.inputTextureDatas[i];
|
|
if (input.isPacked && !artifact.programInfo.expectPackedInputs) {
|
|
runData.inputTextureDatas[i] = this.unpack(input);
|
|
} else if (!input.isPacked && artifact.programInfo.expectPackedInputs) {
|
|
runData.inputTextureDatas[i] = this.pack(input);
|
|
}
|
|
}
|
|
}
|
|
runProgram(artifact: Artifact, runData: RunData) {
|
|
this.checkAndUpdateTextureForm(artifact, runData);
|
|
|
|
// output should match
|
|
if (!!runData.outputTextureData.isPacked !== !!artifact.programInfo.expectPackedOutputs) {
|
|
throw new Error('output property packed inconsistent');
|
|
}
|
|
|
|
this.session.programManager.run(artifact, runData);
|
|
}
|
|
|
|
/**
|
|
* Create a TextureData object from a tensor.
|
|
* Usage = Encoder.Usage.UploadOnly.
|
|
* If a related texture data is found in cache, returns it;
|
|
* Otherwise:
|
|
* Creates a new texture layout if not provided;
|
|
* Creates WebGLTexture with the layout;
|
|
* Upload tensor data to the texture;
|
|
* Creates a texture data object associated with the given tensor.
|
|
* @param tensor the tensor with data to upload
|
|
*/
|
|
getOrCreateTextureData(tensor: Tensor, layout?: TextureLayout, isPacked = false) {
|
|
let td = this.getTextureData(tensor.dataId, isPacked);
|
|
if (!td) {
|
|
Logger.verbose('InferenceHandler', `Creating new TextureData for dims: [${tensor.dims}]`);
|
|
if (!layout) {
|
|
layout = this.createTextureLayoutFromShape(tensor.dims.slice());
|
|
}
|
|
// if we don't find the texture data with specific pack mode in the cache, try with the different
|
|
// pack mode to see if the tensor is cached using that pack mode. If succeed, we can return this
|
|
// tensor data and later apply a pack/unpack op on this texture, no need to create a new one here.
|
|
td = this.getTextureData(tensor.dataId, !isPacked);
|
|
if (!td) {
|
|
if (isPacked) {
|
|
const unpackedTextureLayout = this.getOrCreateTextureLayout(tensor, 1, false, [], true);
|
|
const unpackedTextureData = this.createTextureData(
|
|
unpackedTextureLayout, tensor.type, tensor.numberData, tensor, Encoder.Usage.UploadOnly);
|
|
td = this.pack(unpackedTextureData);
|
|
} else {
|
|
td = this.createTextureData(layout, tensor.type, tensor.numberData, tensor, Encoder.Usage.UploadOnly);
|
|
}
|
|
}
|
|
} else {
|
|
Logger.verbose('InferenceHandler', `Retrieving TextureData from cache: [${tensor.dims}]`);
|
|
}
|
|
return td;
|
|
}
|
|
|
|
/**
|
|
* Create a TextureData object from the given data type and texture layout.
|
|
* Usage = Encoder.Usage.Default.
|
|
* @param dataType the tensor data type
|
|
*/
|
|
createTextureDataFromLayout(layout: TextureLayout, dataType: Tensor.DataType): TextureData {
|
|
return this.createTextureData(layout, dataType);
|
|
}
|
|
|
|
/**
|
|
* Create a TextureData object using the given data and bind to the given tensor.
|
|
* Usage = Encoder.Usage.UploadOnly.
|
|
* NOTE: this function is a hack for Conv implementation. should remove this function, after rewriting Conv
|
|
* implementation by Graph.Transformer
|
|
* @param dataType the tensor data type
|
|
* @param data the actual data to upload
|
|
* @param tensor the tensor to bind. tensor's data is ignored.
|
|
*/
|
|
createTextureDataFromLayoutBindTensor(
|
|
layout: TextureLayout, dataType: Tensor.DataType, data: Tensor.NumberType, tensor: Tensor): TextureData {
|
|
return this.createTextureData(layout, dataType, data, tensor, Encoder.Usage.UploadOnly);
|
|
}
|
|
|
|
private createTextureData(
|
|
layout: TextureLayout, dataType: Tensor.DataType, data?: Tensor.NumberType, tensor?: Tensor,
|
|
usage?: Encoder.Usage): TextureData {
|
|
Logger.verbose('InferenceHandler', `Creating TextureData: layout:[${JSON.stringify(layout)}]`);
|
|
const texture = this.session.textureManager.createTextureFromLayout(dataType, layout, data, usage);
|
|
return this.createTextureDataFromTexture(layout, dataType, texture, tensor);
|
|
}
|
|
|
|
/**
|
|
* Create a TextureData object, using the given texture.
|
|
* This function does not create new texture. Usually used in scenarios using texture sharing. (eg. Reshape)
|
|
* @param dataType the tensor data type
|
|
* @param texture the WebGLTexture object to share
|
|
* @param tensorId the tensor ID of the shared tensor data
|
|
*/
|
|
createSharedTextureData(
|
|
layout: TextureLayout, dataType: Tensor.DataType, texture: WebGLTexture, tensorId?: Tensor.Id): TextureData {
|
|
return this.createTextureDataFromTexture(layout, dataType, texture, undefined, tensorId);
|
|
}
|
|
|
|
private createTextureDataFromTexture(
|
|
layout: TextureLayout, dataType: Tensor.DataType, texture: WebGLTexture, tensor?: Tensor, tensorId?: Tensor.Id) {
|
|
const textureData: TextureData = {
|
|
...layout,
|
|
tensor: tensor ||
|
|
new Tensor(
|
|
layout.unpackedShape, dataType, (_id: Tensor.Id) => this.readTexture(textureData), undefined,
|
|
undefined, tensorId),
|
|
texture
|
|
};
|
|
this.setTextureData(textureData.tensor.dataId, textureData, layout.isPacked);
|
|
return textureData;
|
|
}
|
|
|
|
getTextureData(tensorId: Tensor.Id, isPacked = false): TextureData|undefined {
|
|
return this.session.isInitializer(tensorId) ?
|
|
this.session.getTextureData(tensorId, isPacked) :
|
|
isPacked ? this.packedTextureDataCache.get(tensorId) : this.unpackedTextureDataCache.get(tensorId);
|
|
}
|
|
setTextureData(tensorId: Tensor.Id, td: TextureData, isPacked = false): void {
|
|
if (this.session.isInitializer(tensorId)) {
|
|
this.session.setTextureData(tensorId, td, isPacked);
|
|
} else {
|
|
(isPacked ? this.packedTextureDataCache : this.unpackedTextureDataCache).set(tensorId, td);
|
|
}
|
|
}
|
|
isTextureLayoutCached(tensor: Tensor, isPacked = false): boolean {
|
|
return !!this.getTextureData(tensor.dataId, isPacked);
|
|
}
|
|
/**
|
|
* Create a TextureLayout object from a tensor. If a related texture data is found, returns the cached texture layout.
|
|
*/
|
|
getOrCreateTextureLayout(
|
|
tensor: Tensor, channels: 1|4 = 1, isPacked = false, unpackedShape?: readonly number[],
|
|
reverseWH = false): TextureLayout {
|
|
const td = this.getTextureData(tensor.dataId, isPacked);
|
|
if (td) {
|
|
return td;
|
|
}
|
|
return this.createTextureLayoutFromShape(
|
|
channels === 1 || isPacked ? tensor.dims : getPackedShape(tensor.dims), channels, unpackedShape,
|
|
isPacked || reverseWH ? {isPacked, reverseWH} : undefined);
|
|
}
|
|
|
|
/**
|
|
* Create a TextureLayout object from shape.
|
|
*/
|
|
createTextureLayoutFromShape(
|
|
shape: readonly number[], channels: 1|4 = 1, unpackedShape?: readonly number[],
|
|
prefs?: WidthHeightPrefs): TextureLayout {
|
|
const isPacked = !!(prefs && prefs.isPacked);
|
|
const [texWidth, texHeight] =
|
|
this.session.layoutStrategy.computeTextureWH(isPacked ? unpackedShape || shape : shape, prefs);
|
|
let [width, height] = [texWidth, texHeight];
|
|
if (prefs && prefs.reverseWH) {
|
|
width = texHeight;
|
|
height = texWidth;
|
|
}
|
|
const rank = shape.length;
|
|
let inferredDims = shape.slice(0);
|
|
if (rank === 0) {
|
|
inferredDims = [1];
|
|
}
|
|
if (channels === 1) {
|
|
// unpackedShape will take `shape` and not `inferredDims` so as to create a scalar Tensor if need be
|
|
unpackedShape = shape;
|
|
} else if (isPacked) {
|
|
if (channels !== 4) {
|
|
throw new Error('a packed texture must be 4-channel');
|
|
}
|
|
unpackedShape = shape;
|
|
if (rank > 0) {
|
|
inferredDims[rank - 1] = Math.ceil(inferredDims[rank - 1] / 2);
|
|
}
|
|
if (rank > 1) {
|
|
inferredDims[rank - 2] = Math.ceil(inferredDims[rank - 2] / 2);
|
|
}
|
|
} else if (!unpackedShape) {
|
|
throw new Error('Unpacked shape is needed when using channels > 1');
|
|
}
|
|
return {
|
|
width,
|
|
height,
|
|
channels,
|
|
isPacked,
|
|
shape: inferredDims,
|
|
strides: ShapeUtil.computeStrides(inferredDims),
|
|
unpackedShape,
|
|
reversedWH: (prefs && prefs.reverseWH)
|
|
};
|
|
}
|
|
|
|
dispose(): void {
|
|
this.session.textureManager.clearActiveTextures();
|
|
this.packedTextureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
|
|
this.packedTextureDataCache = new Map();
|
|
this.unpackedTextureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
|
|
this.unpackedTextureDataCache = new Map();
|
|
}
|
|
|
|
readTexture(textureData: TextureData): Tensor.NumberType {
|
|
if (textureData.isPacked) {
|
|
return this.readTexture(this.unpack(textureData));
|
|
}
|
|
if (!this.session.backend.glContext.isFloat32DownloadSupported) {
|
|
const op = new WebGLUint8Encode();
|
|
const uint8TD = op.runInternal(this, textureData);
|
|
return this.session.textureManager.readUint8TextureAsFloat(uint8TD);
|
|
}
|
|
return this.session.textureManager.readTexture(textureData, textureData.tensor.type, textureData.channels);
|
|
}
|
|
|
|
pack(input: TextureData): TextureData {
|
|
const cachedId = this.unpack2packMap.get(input.tensor.dataId);
|
|
if (cachedId) {
|
|
return this.packedTextureDataCache.get(cachedId)!;
|
|
}
|
|
const key = `${input.shape}`;
|
|
let op = this.session.packOpCache.get(key);
|
|
if (!op) {
|
|
op = new WebGLPack();
|
|
this.session.packOpCache.set(key, op);
|
|
}
|
|
let artifact = this.session.programManager.getArtifact(op);
|
|
if (!artifact) {
|
|
const programInfo = op.createProgramInfo(this, [input.tensor]);
|
|
artifact = this.session.programManager.build(programInfo);
|
|
this.session.programManager.setArtifact(op, artifact);
|
|
}
|
|
const runData = op.createRunData(this, artifact.programInfo, [input.tensor]);
|
|
this.runProgram(artifact, runData);
|
|
this.unpack2packMap.set(input.tensor.dataId, runData.outputTextureData.tensor.dataId);
|
|
return runData.outputTextureData;
|
|
}
|
|
|
|
unpack(input: TextureData): TextureData {
|
|
const cachedId = this.pack2unpackMap.get(input.tensor.dataId);
|
|
if (cachedId) {
|
|
return this.unpackedTextureDataCache.get(cachedId)!;
|
|
}
|
|
// For unpacked kernel, cache it by using input's unpackedShape as cache key.
|
|
// Note that we need to use input.unpackedShape instead of input.shape here,
|
|
// as the shape infers the packed texture shape. Different unpackedShape can have the
|
|
// same packed texture shape. For example, for unpacked shape, both [2, 3] and
|
|
// [2, 4] has the same packed shape [1, 2], but those two shapes should have different
|
|
// unpack shaders.
|
|
const key = `${input.unpackedShape}`;
|
|
let op = this.session.unpackOpCache.get(key);
|
|
if (!op) {
|
|
op = new WebGLUnpack();
|
|
this.session.unpackOpCache.set(key, op);
|
|
}
|
|
let artifact = this.session.programManager.getArtifact(op);
|
|
if (!artifact) {
|
|
const programInfo = op.createProgramInfo(this, [input.tensor]);
|
|
artifact = this.session.programManager.build(programInfo);
|
|
this.session.programManager.setArtifact(op, artifact);
|
|
}
|
|
const runData = op.createRunData(this, artifact.programInfo, [input.tensor]);
|
|
this.runProgram(artifact, runData);
|
|
this.pack2unpackMap.set(input.tensor.dataId, runData.outputTextureData.tensor.dataId);
|
|
return runData.outputTextureData;
|
|
}
|
|
}
|