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

Commit

Permalink
Merge remote-tracking branch 'origin' into duli/conv_fix
Browse files Browse the repository at this point in the history
  • Loading branch information
@duli2012
duli2012 committed May 4, 2021
2 parents e185132 + 94d3843 commit 6cf4689
Show file tree
Hide file tree
Showing 17 changed files with 225 additions and 120 deletions.
38 changes: 31 additions & 7 deletions lib/backends/webgl/glsl-coordinate-lib.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ export class CoordsGlslLib extends GlslLib {
* Generates code for packed output sampler.
*/
protected getPackedOutputSamplingSnippet(outputLayout: TextureLayout): {[name: string]: GlslLibRoutine;} {
const outShape = outputLayout.shape;
const outShape = outputLayout.unpackedShape;
const outTexShape = [outputLayout.width, outputLayout.height];
const result: {[name: string]: GlslLibRoutine} = {};
const funcName = 'getOutputCoords';
Expand Down Expand Up @@ -231,7 +231,7 @@ export class CoordsGlslLib extends GlslLib {

const packedTexShape = texShape;
// texels needed to accommodate a logical row
const texelsInLogicalRow = shape[1];
const texelsInLogicalRow = Math.ceil(shape[1] / 2);

/**
* getOutputCoords
Expand Down Expand Up @@ -264,8 +264,9 @@ export class CoordsGlslLib extends GlslLib {
*/
protected getOutputPacked3DCoords(shape: [number, number, number], texShape: [number, number]): GlslLibRoutine {
const packedTexShape = [texShape[0], texShape[1]];
const texelsInLogicalRow = shape[2];
const texelsInBatch = texelsInLogicalRow * shape[1];
const texelsInLogicalRow = Math.ceil(shape[2] / 2);

const texelsInBatch = texelsInLogicalRow * Math.ceil(shape[1] / 2);
const source = `
ivec3 getOutputCoords() {
ivec2 resTexRC = ivec2(TexCoords.xy *
Expand All @@ -291,8 +292,8 @@ export class CoordsGlslLib extends GlslLib {
protected getOutputPackedNDCoords(shape: ReadonlyArray<number>, texShape: [number, number]): GlslLibRoutine {
const packedTexShape = [texShape[0], texShape[1]];

const texelsInLogicalRow = shape[shape.length - 1];
const texelsInBatch = texelsInLogicalRow * shape[shape.length - 2];
const texelsInLogicalRow = Math.ceil(shape[shape.length - 1] / 2);
const texelsInBatch = texelsInLogicalRow * Math.ceil(shape[shape.length - 2] / 2);
let texelsInBatchN = texelsInBatch;
let batches = ``;
let coords = 'b, r, c';
Expand Down Expand Up @@ -690,7 +691,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 +1217,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
14 changes: 10 additions & 4 deletions lib/backends/webgl/ops/concat_packed.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ export class WebGLPackedConcat extends Concat implements WebGLOperator {
const unpackChannel = unpackFromChannel();

const shapes = inputs.map(i => i.dims);
const channels = ['x', 'y', 'z', 'w', 'u', 'v'].slice(0, rank);
const allGlChannels = ['x', 'y', 'z', 'w', 'u', 'v'];
const channels = allGlChannels.slice(0, rank);
const offsets: number[] = new Array(shapes.length - 1);
const samplers = inputs.map((v, i) => `X${i}`);

Expand Down Expand Up @@ -88,6 +89,10 @@ export class WebGLPackedConcat extends Concat implements WebGLOperator {
void main() {
${dtype} coords = getOutputCoords();
int lastDim = coords.${allGlChannels[rank - 1]};
coords.${allGlChannels[rank - 1]} = coords.${allGlChannels[rank - 2]};
coords.${allGlChannels[rank - 2]} = lastDim;
vec4 result = vec4(getValue(${coords}), 0., 0., 0.);
${coords[rank - 1]} = ${coords[rank - 1]} + 1;
Expand All @@ -110,8 +115,9 @@ export class WebGLPackedConcat extends Concat implements WebGLOperator {
`;

return {
inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t)),
outputLayout: handler.createTextureLayoutFromShape(outputShape),
inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)),
outputLayout:
handler.createTextureLayoutFromShape(outputShape, 4, outputShape, {isPacked: true, reverseWH: true}),
samplers,
shaderSource,
hasMain: true,
Expand All @@ -120,7 +126,7 @@ export class WebGLPackedConcat extends Concat implements WebGLOperator {
};
}
createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
const inputTDs = inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i]));
const inputTDs = inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i], true));
return {
inputTextureDatas: inputTDs,
outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
Expand Down
Loading

0 comments on commit 6cf4689

Please sign in to comment.