Skip to content
This repository has been archived by the owner on Nov 16, 2023. It is now read-only.

Commit

Permalink
add texture cache support for pack/unpack mode (#290)
Browse files Browse the repository at this point in the history 
* squash change

* uncomment test verification

* clean up

* clean up and remove unused files

* adding comment

* remove unused code

* fix reshape merge

* add guid as tensor id

* add file

* fix test failure
  • Loading branch information
@xzhu1900
xzhu1900 authored Apr 28, 2021
1 parent 290825d commit 0c998cc
Show file tree
Hide file tree
Showing 15 changed files with 192 additions and 77 deletions.
25 changes: 24 additions & 1 deletion lib/backends/webgl/glsl-coordinate-lib.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -690,7 +690,7 @@ export class CoordsGlslLib extends GlslLib {
return ${texFuncSnippet}(${unpackedCoordsSnippet});
}
`;
return new GlslLibRoutine(source);
return new GlslLibRoutine(source, ['coordinates.getOutputCoords']);
}

/**
Expand Down Expand Up @@ -1216,4 +1216,27 @@ export class CoordsGlslLib extends GlslLib {
}
`;
}

/**
* Produces a packed value getter function for the name and rank given
* If a transpose is set proper offsetToCoords mapping will be used
* @param name name of the function
* @param rank rank of the input
* @param transpose whether or not should generate a transpose variation
*/
protected getPackedValueFrom(varName: string, rank: number, width: number, height: number, transpose: boolean):
string {
let name = `_${varName}_Pack`;
if (transpose) {
name = name + '_T';
}
const glsl = getGlsl(this.context.glContext.version);
return `
vec4 ${name}(int m[${rank}]) {
int offset = indicesToOffset_${varName}(m);
vec2 coords = offsetToCoords(offset, ${width}, ${height});
return ${glsl.texture2D}(${varName}, coords);
}
`;
}
}
114 changes: 67 additions & 47 deletions lib/backends/webgl/inference-handler.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,16 @@ import {Artifact, RunData, TextureData, TextureLayout, WebGLOperator} from './ty
import {getPackedShape} from './utils';

export class WebGLInferenceHandler implements InferenceHandler {
private textureDataCache: Map<Tensor.Id, TextureData>;
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.textureDataCache = new Map();
this.packedTextureDataCache = new Map();
this.unpackedTextureDataCache = new Map();

this.pack2unpackMap = new Map();
this.unpack2packMap = new Map();
}

run(op: WebGLOperator, inputs: Tensor[]): Tensor[] {
Expand All @@ -33,36 +40,18 @@ export class WebGLInferenceHandler implements InferenceHandler {
return [runData.outputTextureData.tensor];
}

/**
* Check the runData's input texture mode with the program's artifact.
* If the artifact expects a packed input, while the RunData's input
* is unpacked, perform a pack operation on this input to align the
* texture mode with artifact. Similar on unpacked input.
*/
checkAndUpdateTextureForm(artifact: Artifact, runData: RunData) {
// pack/unpack inputs
runData.inputTextureDatas.forEach(input => {
for (let i = 0; i < runData.inputTextureDatas.length; ++i) {
const input = runData.inputTextureDatas[i];
if (input.isPacked && !artifact.programInfo.expectPackedInputs) {
// unpack this input
const unpacked = this.unpack(input);
input.height = unpacked.height;
input.isPacked = unpacked.isPacked;
input.texture = unpacked.texture;
input.width = unpacked.width;

runData.inputTextureDatas[i] = this.unpack(input);
} else if (!input.isPacked && artifact.programInfo.expectPackedInputs) {
// pack this input
const packed = this.pack(input);
input.height = packed.height;
input.isPacked = packed.isPacked;
input.texture = packed.texture;
input.width = packed.width;
runData.inputTextureDatas[i] = this.pack(input);
}
});
}
}
runProgram(artifact: Artifact, runData: RunData) {
// if the runData has different expected texture pack/unpack mode, process pack/unpack
// operation on the texture before executing the kernel.
this.checkAndUpdateTextureForm(artifact, runData);

// output should match
Expand All @@ -84,15 +73,28 @@ export class WebGLInferenceHandler implements InferenceHandler {
* Creates a texture data object associated with the given tensor.
* @param tensor the tensor with data to upload
*/
getOrCreateTextureData(tensor: Tensor, layout?: TextureLayout) {
let td = this.getTextureData(tensor.dataId);
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());
}
// graph inputs or initializers
td = this.createTextureData(layout, tensor.type, tensor.numberData, tensor, Encoder.Usage.UploadOnly);
// 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, isPacked);
}
}
} else {
Logger.verbose('InferenceHandler', `Retrieving TextureData from cache: [${tensor.dims}]`);
}
Expand All @@ -104,7 +106,7 @@ export class WebGLInferenceHandler implements InferenceHandler {
* Usage = Encoder.Usage.Default.
* @param dataType the tensor data type
*/
createTextureDataFromLayout(layout: TextureLayout, dataType: Tensor.DataType): TextureData {
createTextureDataFromLayout(layout: TextureLayout, dataType: Tensor.DataType, isPacked = false): TextureData {
return this.createTextureData(layout, dataType);
}

Expand All @@ -118,13 +120,14 @@ export class WebGLInferenceHandler implements InferenceHandler {
* @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);
layout: TextureLayout, dataType: Tensor.DataType, data: Tensor.NumberType, tensor: Tensor,
isPacked = false): TextureData {
return this.createTextureData(layout, dataType, data, tensor, Encoder.Usage.UploadOnly, isPacked);
}

private createTextureData(
layout: TextureLayout, dataType: Tensor.DataType, data?: Tensor.NumberType, tensor?: Tensor,
usage?: Encoder.Usage): TextureData {
usage?: Encoder.Usage, isPacked = false): 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);
Expand All @@ -137,8 +140,9 @@ export class WebGLInferenceHandler implements InferenceHandler {
* @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 {
createSharedTextureData(
layout: TextureLayout, dataType: Tensor.DataType, texture: WebGLTexture, tensorId?: Tensor.Id,
isPacked = false): TextureData {
return this.createTextureDataFromTexture(layout, dataType, texture, undefined, tensorId);
}

Expand All @@ -155,29 +159,32 @@ export class WebGLInferenceHandler implements InferenceHandler {
undefined, undefined, tensorId),
texture
};
this.setTextureData(textureData.tensor.dataId, textureData);
this.setTextureData(textureData.tensor.dataId, textureData, layout.isPacked);
return textureData;
}

getTextureData(tensorId: Tensor.Id): TextureData|undefined {
return this.session.isInitializer(tensorId) ? this.session.getTextureData(tensorId) :
this.textureDataCache.get(tensorId);
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): void {
setTextureData(tensorId: Tensor.Id, td: TextureData, isPacked = false): void {
if (this.session.isInitializer(tensorId)) {
this.session.setTextureData(tensorId, td);
this.session.setTextureData(tensorId, td, isPacked);
} else {
this.textureDataCache.set(tensorId, td);
isPacked ? this.packedTextureDataCache.set(tensorId, td) : 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?: ReadonlyArray<number>,
reverseWH = false): TextureLayout {
const td = this.getTextureData(tensor.dataId);
const td = this.getTextureData(tensor.dataId, isPacked);
if (td) {
return td;
}
Expand Down Expand Up @@ -229,14 +236,17 @@ export class WebGLInferenceHandler implements InferenceHandler {
isPacked,
shape: inferredDims,
strides: ShapeUtil.computeStrides(inferredDims),
unpackedShape
unpackedShape,
reversedWH: (prefs && prefs.reverseWH)
};
}

dispose(): void {
this.session.textureManager.clearActiveTextures();
this.textureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
this.textureDataCache = new Map();
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 {
Expand All @@ -252,6 +262,10 @@ export class WebGLInferenceHandler implements InferenceHandler {
}

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) {
Expand All @@ -266,10 +280,15 @@ export class WebGLInferenceHandler implements InferenceHandler {
}
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
Expand All @@ -290,6 +309,7 @@ export class WebGLInferenceHandler implements InferenceHandler {
}
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;
}
}
10 changes: 9 additions & 1 deletion lib/backends/webgl/ops/binary-op.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
const inputLayouts = handler.session.pack ?
inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)) :
inputs.map(t => handler.getOrCreateTextureLayout(t));
const ouputLayout = handler.session.pack ?
handler.createTextureLayoutFromShape(inputs[0].dims, 4, inputs[0].dims, {isPacked: true, reverseWH: true}) :
handler.createTextureLayoutFromShape(inputs[0].dims);

const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
if (isBroadcast) {
const outputShape = BroadcastUtil.calcShape(inputs[0].dims, inputs[1].dims, false);
Expand All @@ -33,6 +37,8 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
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 = `
${this.glslFunc.body}
float process(int indices[${outputRank}]) {
Expand All @@ -51,6 +57,8 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
outputLayout,
samplers: ['A', 'B'],
shaderSource,
expectPackedInputs: handler.session.pack,
expectPackedOutputs: handler.session.pack
};
}
const glsl = getGlsl(handler.session.backend.glContext.version);
Expand All @@ -67,7 +75,7 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
return {
hasMain: true,
inputLayouts,
outputLayout: handler.createTextureLayoutFromShape(inputs[0].dims),
outputLayout: ouputLayout,
samplers: ['A', 'B'],
shaderSource,
expectPackedInputs: true,
Expand Down
9 changes: 8 additions & 1 deletion lib/backends/webgl/ops/reshape-packed.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,10 +117,17 @@ export class WebGLReshapePacked extends Reshape implements WebGLOperator {
createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
const inputTDs =
[handler.getOrCreateTextureData(inputs[0], handler.getOrCreateTextureLayout(inputs[0], 1, false, [], false))];
let outputLayout = this.originalOutputLayout;
if (outputLayout === undefined) {
const originInputShape = inputs[0].dims;
const outputShape = ShapeUtil.calculateReshapedDims(originInputShape, inputs[1].integerData);
outputLayout =
handler.createTextureLayoutFromShape(outputShape, 4, outputShape, {isPacked: true, reverseWH: true});
}
// return run data for reshape. Here, we use the original calculate outputLayout to create the real output layout.
return {
inputTextureDatas: inputTDs,
outputTextureData: handler.createTextureDataFromLayout(this.originalOutputLayout, inputTDs[0].tensor.type),
outputTextureData: handler.createTextureDataFromLayout(outputLayout, inputTDs[0].tensor.type),
uniformData: {}
};
}
Expand Down
2 changes: 1 addition & 1 deletion lib/backends/webgl/ops/reshape.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,6 @@ export function reshape(
unpackedShape: reshapedDims,
};

const newTextureData = inferenceHandler.createSharedTextureData(newTextureLayout, input.type, inputTD.texture, {});
const newTextureData = inferenceHandler.createSharedTextureData(newTextureLayout, input.type, inputTD.texture);
return newTextureData.tensor;
}
2 changes: 1 addition & 1 deletion lib/backends/webgl/ops/uint8-encode.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ export class WebGLUint8Encode {
const encoder = inferenceHandler.session.backend.glContext.getEncoder('byte', 4);
const texture =
inferenceHandler.session.backend.glContext.allocateTexture(outputLayout.width, outputLayout.height, encoder);
const outputTextureData = inferenceHandler.createSharedTextureData(outputLayout, 'uint8', texture, {});
const outputTextureData = inferenceHandler.createSharedTextureData(outputLayout, 'uint8', texture);
const runData = {inputTextureDatas: [input], outputTextureData, uniformData: {}};

inferenceHandler.session.programManager.run(artifact, runData);
Expand Down
7 changes: 3 additions & 4 deletions lib/backends/webgl/ops/unpack.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ import {getGlsl} from '../glsl-source';
import {WebGLInferenceHandler} from '../inference-handler';
import {ProgramInfo, RunData, WebGLOperator} from '../types';
import {getCoordsDataType} from '../utils';

import {getChannels, unpackFromChannel} from './packing_utils';

export class WebGLUnpack implements WebGLOperator {
Expand All @@ -18,7 +17,7 @@ export class WebGLUnpack implements WebGLOperator {
throw new Error(`Pack kernel should have input tensor count to 1.`);
}

const inputTexture = handler.getTextureData(inputs[0].dataId);
const inputTexture = handler.getTextureData(inputs[0].dataId, true);
if (!inputTexture) {
throw new Error(`packed input texture must exist`);
}
Expand Down Expand Up @@ -49,7 +48,7 @@ export class WebGLUnpack implements WebGLOperator {
`;

return {
inputLayouts: [handler.getOrCreateTextureLayout(inputs[0])],
inputLayouts: [handler.getOrCreateTextureLayout(inputs[0], 4, true, inputs[0].dims, true)],
outputLayout,
samplers: ['A'],
shaderSource,
Expand All @@ -59,7 +58,7 @@ export class WebGLUnpack implements WebGLOperator {
};
}
createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
const inputTDs = [handler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0])];
const inputTDs = [handler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0], true)];
return {
inputTextureDatas: inputTDs,
outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
Expand Down
7 changes: 4 additions & 3 deletions lib/backends/webgl/program-manager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,11 +105,12 @@ ${fragShaderScript}
return program;
}
bindOutput(td: TextureData): void {
const width = td.width;
const height = td.height;
Logger.verbose(
'ProrgramManager',
`Binding output texture to Framebuffer: w/h=${td.width}/${td.height}, shape=${td.shape}, type=${
td.tensor.type}`);
this.glContext.attachFramebuffer(td.texture, td.width, td.height);
`Binding output texture to Framebuffer: w/h=${width}/${height}, shape=${td.shape}, type=${td.tensor.type}`);
this.glContext.attachFramebuffer(td.texture, width, height);
}
bindAttributes(attribLocations: Artifact.AttribLocations): void {
const positionHandle = attribLocations.position;
Expand Down
Loading

0 comments on commit 0c998cc

Please sign in to comment.