From 2b9bc1c0faa14286af4217aea8718b52da59cc53 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 16 Feb 2023 13:05:20 -0800 Subject: [PATCH 1/8] Improve texture manager array performance --- tfjs-backend-webgl/src/texture_manager.ts | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tfjs-backend-webgl/src/texture_manager.ts b/tfjs-backend-webgl/src/texture_manager.ts index 1d2d174c6a9..24c58f9ce62 100644 --- a/tfjs-backend-webgl/src/texture_manager.ts +++ b/tfjs-backend-webgl/src/texture_manager.ts @@ -27,11 +27,11 @@ export class TextureManager { private _numBytesAllocated = 0; private _numBytesFree = 0; // How many bytes that have been allocated // are available for reuse. - private freeTextures: {[shape: string]: Texture[]} = {}; + private freeTextures: Record = {}; + private usedTextures: Record = {}; private logEnabled = false; - private usedTextures: {[shape: string]: Texture[]} = {}; - constructor(private gpgpu: GPGPUContext) {} + constructor(private readonly gpgpu: GPGPUContext) {} acquireTexture( shapeRC: [number, number], usage: TextureUsage, @@ -55,7 +55,7 @@ export class TextureManager { this.numUsedTextures++; this._numBytesFree -= texBytes; this.log(); - const newTexture = this.freeTextures[shapeKey].shift(); + const newTexture = this.freeTextures[shapeKey].pop(); this.usedTextures[shapeKey].push(newTexture); return newTexture; } @@ -123,7 +123,8 @@ export class TextureManager { 'Cannot release a texture that was never provided by this ' + 'texture manager'); } - texList.splice(texIndex, 1); + texList[texIndex] = texList[texList.length - 1]; + texList.pop(); this.log(); } @@ -172,8 +173,8 @@ export class TextureManager { this.gpgpu.deleteMatrixTexture(tex.texture); }); } - this.freeTextures = null; - this.usedTextures = null; + this.freeTextures = {}; + this.usedTextures = {}; this.numUsedTextures = 0; this.numFreeTextures = 0; this._numBytesAllocated = 0; From 02dabf5e38588fbf89e482c113490e593a5f9bb0 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 16 Feb 2023 13:40:02 -0800 Subject: [PATCH 2/8] Improve uniform loc lookup performance --- tfjs-backend-webgl/src/backend_webgl.ts | 8 +-- tfjs-backend-webgl/src/gpgpu_math.ts | 69 +++++++++++------------ tfjs-backend-webgl/src/texture_manager.ts | 4 +- 3 files changed, 36 insertions(+), 45 deletions(-) diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index e3b8970cf5a..5da5b8dbc2e 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -1278,22 +1278,18 @@ export class MathBackendWebGL extends KernelBackend { getUniformLocations() { for (const [, binary] of Object.entries(this.binaryCache)) { const { - uniformLocations, + variableLocations, customUniformLocations, infLoc, nanLoc, - inShapesLocations, - inTexShapesLocations, outShapeLocation, outShapeStridesLocation, outTexShapeLocation } = getUniformLocations(this.gpgpu, binary.program, binary.webGLProgram); - binary.uniformLocations = uniformLocations; + binary.variableLocations = variableLocations; binary.customUniformLocations = customUniformLocations; binary.infLoc = infLoc; binary.nanLoc = nanLoc; - binary.inShapesLocations = inShapesLocations; - binary.inTexShapesLocations = inTexShapesLocations; binary.outShapeLocation = outShapeLocation; binary.outShapeStridesLocation = outShapeStridesLocation; binary.outTexShapeLocation = outTexShapeLocation; diff --git a/tfjs-backend-webgl/src/gpgpu_math.ts b/tfjs-backend-webgl/src/gpgpu_math.ts index 4ee65d29c99..8b3906d5598 100644 --- a/tfjs-backend-webgl/src/gpgpu_math.ts +++ b/tfjs-backend-webgl/src/gpgpu_math.ts @@ -46,34 +46,31 @@ export interface GPGPUProgram { Array<{name: string; arrayIndex?: number; type: UniformType;}>; } -export interface GPGPUBinary { +export interface GPGPUBinary extends GPGPUBinaryLocations { webGLProgram: GPGPUContextProgram; program: GPGPUProgram; - uniformLocations: {[name: string]: WebGLUniformLocation}; - customUniformLocations?: WebGLUniformLocation[]; source: string; fragmentShader: WebGLShader; inShapeInfos: ShapeInfo[]; outShapeInfo: ShapeInfo; - infLoc: WebGLUniformLocation; - nanLoc: WebGLUniformLocation; - inShapesLocations?: {[name: string]: WebGLUniformLocation}; - inTexShapesLocations?: {[name: string]: WebGLUniformLocation}; - outShapeLocation?: WebGLUniformLocation; - outShapeStridesLocation?: WebGLUniformLocation; - outTexShapeLocation?: WebGLUniformLocation; } export interface GPGPUBinaryLocations { - uniformLocations: {[name: string]: WebGLUniformLocation}; customUniformLocations?: WebGLUniformLocation[]; infLoc: WebGLUniformLocation; nanLoc: WebGLUniformLocation; - inShapesLocations?: {[name: string]: WebGLUniformLocation}; - inTexShapesLocations?: {[name: string]: WebGLUniformLocation}; outShapeLocation?: WebGLUniformLocation; outShapeStridesLocation?: WebGLUniformLocation; outTexShapeLocation?: WebGLUniformLocation; + variableLocations?: GPGPUVariableLocations[]; +} + +export interface GPGPUVariableLocations { + variableName: string; + uniform: WebGLUniformLocation; + offset: WebGLUniformLocation; + shape?: WebGLUniformLocation; + texShape?: WebGLUniformLocation; } export interface TensorData { @@ -131,12 +128,10 @@ export function compileProgram( webGLProgram, inShapeInfos, outShapeInfo, - uniformLocations: null, + variableLocations: [], customUniformLocations: null, infLoc: null, nanLoc: null, - inShapesLocations: null, - inTexShapesLocations: null, outShapeLocation: null, outShapeStridesLocation: null, outTexShapeLocation: null @@ -147,9 +142,7 @@ export function compileProgram( export function getUniformLocations( gpgpu: GPGPUContext, program: GPGPUProgram, webGLProgram: WebGLProgram): GPGPUBinaryLocations { - const uniformLocations: {[name: string]: WebGLUniformLocation} = {}; - const inShapesLocations: {[name: string]: WebGLUniformLocation} = {}; - const inTexShapesLocations: {[name: string]: WebGLUniformLocation} = {}; + const variableLocations: GPGPUVariableLocations[] = []; const customUniformLocations: WebGLUniformLocation[] = []; let outShapeLocation: WebGLUniformLocation; let outTexShapeLocation: WebGLUniformLocation; @@ -165,18 +158,20 @@ export function getUniformLocations( // Add user-defined uniforms const shouldThrow = false; - for (let i = 0; i < program.variableNames.length; i++) { - const varName = program.variableNames[i]; - uniformLocations[varName] = - gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow); - uniformLocations[`offset${varName}`] = - gpgpu.getUniformLocation(webGLProgram, `offset${varName}`, shouldThrow); + for (const varName of program.variableNames) { + let loc: GPGPUVariableLocations = { + variableName: varName, + uniform: gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow), + offset: gpgpu.getUniformLocation( + webGLProgram, `offset${varName}`, shouldThrow), + }; if (program.enableShapeUniforms) { - inShapesLocations[`${varName}Shape`] = gpgpu.getUniformLocation( + loc.shape = gpgpu.getUniformLocation( webGLProgram, `${varName}Shape`, shouldThrow); - inTexShapesLocations[`${varName}TexShape`] = gpgpu.getUniformLocation( + loc.texShape = gpgpu.getUniformLocation( webGLProgram, `${varName}TexShape`, shouldThrow); } + variableLocations.push(loc); } if (program.enableShapeUniforms) { @@ -196,12 +191,10 @@ export function getUniformLocations( } return { - uniformLocations, + variableLocations, customUniformLocations, infLoc, nanLoc, - inShapesLocations, - inTexShapesLocations, outShapeLocation, outShapeStridesLocation, outTexShapeLocation @@ -271,12 +264,14 @@ export function runProgram( } // Set user-defined inputs - inputs.forEach((input, i) => { - const varName = binary.program.variableNames[i]; - const varLoc = binary.uniformLocations[varName]; - const varOffsetLoc = binary.uniformLocations[`offset${varName}`]; - const varShapeLoc = binary.inShapesLocations[`${varName}Shape`]; - const varTexShapeLoc = binary.inTexShapesLocations[`${varName}TexShape`]; + for (let i = 0; i < inputs.length; ++i) { + const input = inputs[i]; + const { + uniform: varLoc, + offset: varOffsetLoc, + shape: varShapeLoc, + texShape: varTexShapeLoc, + } = binary.variableLocations[i]; if (varShapeLoc) { const {uniformShape} = shader_compiler.getUniformInfoFromShape( @@ -328,7 +323,7 @@ export function runProgram( } gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i); - }); + } const outShapeLoc = binary.outShapeLocation; if (outShapeLoc) { diff --git a/tfjs-backend-webgl/src/texture_manager.ts b/tfjs-backend-webgl/src/texture_manager.ts index 24c58f9ce62..1494d702ec0 100644 --- a/tfjs-backend-webgl/src/texture_manager.ts +++ b/tfjs-backend-webgl/src/texture_manager.ts @@ -117,8 +117,8 @@ export class TextureManager { this.numUsedTextures--; const texList = this.usedTextures[shapeKey]; - const texIndex = texList.indexOf(texture); - if (texIndex < 0) { + const texIndex = texList && texList.indexOf(texture); + if (texIndex == null || texIndex < 0) { throw new Error( 'Cannot release a texture that was never provided by this ' + 'texture manager'); From 548a10c3ed84d2eb8c5d964349fee11c3b6429f0 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 16 Feb 2023 14:13:15 -0800 Subject: [PATCH 3/8] Fix texture manager test --- tfjs-backend-webgl/src/texture_manager.ts | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tfjs-backend-webgl/src/texture_manager.ts b/tfjs-backend-webgl/src/texture_manager.ts index 1494d702ec0..8173aa3caad 100644 --- a/tfjs-backend-webgl/src/texture_manager.ts +++ b/tfjs-backend-webgl/src/texture_manager.ts @@ -25,10 +25,10 @@ export class TextureManager { private numUsedTextures = 0; private numFreeTextures = 0; private _numBytesAllocated = 0; - private _numBytesFree = 0; // How many bytes that have been allocated - // are available for reuse. - private freeTextures: Record = {}; - private usedTextures: Record = {}; + // Number of bytes that have been allocated and available for reuse. + private _numBytesFree = 0; + private freeTextures?: Record = {}; + private usedTextures?: Record = {}; private logEnabled = false; constructor(private readonly gpgpu: GPGPUContext) {} @@ -173,8 +173,8 @@ export class TextureManager { this.gpgpu.deleteMatrixTexture(tex.texture); }); } - this.freeTextures = {}; - this.usedTextures = {}; + this.freeTextures = null; + this.usedTextures = null; this.numUsedTextures = 0; this.numFreeTextures = 0; this._numBytesAllocated = 0; From d89b1d63f2012a4cfbcd97380c47f2f8eb84a8a3 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 16 Feb 2023 14:51:36 -0800 Subject: [PATCH 4/8] fix --- tfjs-backend-webgl/src/backend_webgl.ts | 2 +- tfjs-backend-webgl/src/gpgpu_math.ts | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index 5da5b8dbc2e..d4a06db395e 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -1276,7 +1276,7 @@ export class MathBackendWebGL extends KernelBackend { } getUniformLocations() { - for (const [, binary] of Object.entries(this.binaryCache)) { + for (const binary of Object.values(this.binaryCache)) { const { variableLocations, customUniformLocations, diff --git a/tfjs-backend-webgl/src/gpgpu_math.ts b/tfjs-backend-webgl/src/gpgpu_math.ts index 8b3906d5598..8d9deaae5b2 100644 --- a/tfjs-backend-webgl/src/gpgpu_math.ts +++ b/tfjs-backend-webgl/src/gpgpu_math.ts @@ -128,7 +128,7 @@ export function compileProgram( webGLProgram, inShapeInfos, outShapeInfo, - variableLocations: [], + variableLocations: null, customUniformLocations: null, infLoc: null, nanLoc: null, @@ -159,7 +159,7 @@ export function getUniformLocations( // Add user-defined uniforms const shouldThrow = false; for (const varName of program.variableNames) { - let loc: GPGPUVariableLocations = { + const loc: GPGPUVariableLocations = { variableName: varName, uniform: gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow), offset: gpgpu.getUniformLocation( @@ -171,6 +171,7 @@ export function getUniformLocations( loc.texShape = gpgpu.getUniformLocation( webGLProgram, `${varName}TexShape`, shouldThrow); } + variableLocations.push(loc); } @@ -370,7 +371,8 @@ export function runProgram( } if (binary.program.customUniforms && customUniformValues) { - binary.program.customUniforms.forEach((d, i) => { + for (let i = 0; i < binary.program.customUniforms.length; ++i) { + const d = binary.program.customUniforms[i]; const customLoc = binary.customUniformLocations[i]; const customValue = customUniformValues[i]; if (d.type === 'float') { @@ -392,7 +394,7 @@ export function runProgram( } else { throw Error(`uniform type ${d.type} is not supported yet.`); } - }); + } } gpgpu.executeProgram(); } From 987a7691bded7516bc76fc9179c0f07fa210e85d Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Thu, 16 Feb 2023 15:33:52 -0800 Subject: [PATCH 5/8] fix --- tfjs-backend-webgl/src/backend_webgl.ts | 4 +-- tfjs-backend-webgl/src/gpgpu_math.ts | 37 +++++++++++++------------ 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index d4a06db395e..7804901be77 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -1278,7 +1278,7 @@ export class MathBackendWebGL extends KernelBackend { getUniformLocations() { for (const binary of Object.values(this.binaryCache)) { const { - variableLocations, + variablesLocations, customUniformLocations, infLoc, nanLoc, @@ -1286,7 +1286,7 @@ export class MathBackendWebGL extends KernelBackend { outShapeStridesLocation, outTexShapeLocation } = getUniformLocations(this.gpgpu, binary.program, binary.webGLProgram); - binary.variableLocations = variableLocations; + binary.variablesLocations = variablesLocations; binary.customUniformLocations = customUniformLocations; binary.infLoc = infLoc; binary.nanLoc = nanLoc; diff --git a/tfjs-backend-webgl/src/gpgpu_math.ts b/tfjs-backend-webgl/src/gpgpu_math.ts index 8d9deaae5b2..1ca84e8d514 100644 --- a/tfjs-backend-webgl/src/gpgpu_math.ts +++ b/tfjs-backend-webgl/src/gpgpu_math.ts @@ -62,11 +62,11 @@ export interface GPGPUBinaryLocations { outShapeLocation?: WebGLUniformLocation; outShapeStridesLocation?: WebGLUniformLocation; outTexShapeLocation?: WebGLUniformLocation; - variableLocations?: GPGPUVariableLocations[]; + variablesLocations?: GPGPUVariableLocations[]; } export interface GPGPUVariableLocations { - variableName: string; + name: string; uniform: WebGLUniformLocation; offset: WebGLUniformLocation; shape?: WebGLUniformLocation; @@ -128,7 +128,7 @@ export function compileProgram( webGLProgram, inShapeInfos, outShapeInfo, - variableLocations: null, + variablesLocations: null, customUniformLocations: null, infLoc: null, nanLoc: null, @@ -142,7 +142,7 @@ export function compileProgram( export function getUniformLocations( gpgpu: GPGPUContext, program: GPGPUProgram, webGLProgram: WebGLProgram): GPGPUBinaryLocations { - const variableLocations: GPGPUVariableLocations[] = []; + const variablesLocations: GPGPUVariableLocations[] = []; const customUniformLocations: WebGLUniformLocation[] = []; let outShapeLocation: WebGLUniformLocation; let outTexShapeLocation: WebGLUniformLocation; @@ -159,20 +159,20 @@ export function getUniformLocations( // Add user-defined uniforms const shouldThrow = false; for (const varName of program.variableNames) { - const loc: GPGPUVariableLocations = { - variableName: varName, + const varLocs: GPGPUVariableLocations = { + name: varName, uniform: gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow), offset: gpgpu.getUniformLocation( webGLProgram, `offset${varName}`, shouldThrow), }; if (program.enableShapeUniforms) { - loc.shape = gpgpu.getUniformLocation( + varLocs.shape = gpgpu.getUniformLocation( webGLProgram, `${varName}Shape`, shouldThrow); - loc.texShape = gpgpu.getUniformLocation( + varLocs.texShape = gpgpu.getUniformLocation( webGLProgram, `${varName}TexShape`, shouldThrow); } - variableLocations.push(loc); + variablesLocations.push(varLocs); } if (program.enableShapeUniforms) { @@ -185,14 +185,14 @@ export function getUniformLocations( } if (program.customUniforms) { - program.customUniforms.forEach((d, i) => { - customUniformLocations[i] = - gpgpu.getUniformLocation(webGLProgram, d.name, shouldThrow); - }); + for (const d of program.customUniforms) { + customUniformLocations.push( + gpgpu.getUniformLocation(webGLProgram, d.name, shouldThrow)); + } } return { - variableLocations, + variablesLocations, customUniformLocations, infLoc, nanLoc, @@ -272,7 +272,7 @@ export function runProgram( offset: varOffsetLoc, shape: varShapeLoc, texShape: varTexShapeLoc, - } = binary.variableLocations[i]; + } = binary.variablesLocations[i]; if (varShapeLoc) { const {uniformShape} = shader_compiler.getUniformInfoFromShape( @@ -294,6 +294,7 @@ export function runProgram( break; } } + if (varTexShapeLoc) { gpgpu.gl.uniform2i( varTexShapeLoc, input.texData.texShape[0], input.texData.texShape[1]); @@ -301,7 +302,7 @@ export function runProgram( if (varLoc == null) { // The compiler inferred that this variable is not used in this shader. - return; + continue; } if (input.isUniform) { @@ -315,7 +316,7 @@ export function runProgram( } gpgpu.gl.uniform1fv(varLoc, vals); } - return; + continue; } // If the input was sliced, upload the flat offset index. @@ -324,7 +325,7 @@ export function runProgram( } gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i); - } + }; const outShapeLoc = binary.outShapeLocation; if (outShapeLoc) { From 94deac6641758f8577b066421ad3b45b6d3c86e5 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Fri, 17 Feb 2023 10:11:32 -0800 Subject: [PATCH 6/8] Fix lint --- tfjs-backend-webgl/src/gpgpu_math.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgl/src/gpgpu_math.ts b/tfjs-backend-webgl/src/gpgpu_math.ts index 1ca84e8d514..9de19971ad1 100644 --- a/tfjs-backend-webgl/src/gpgpu_math.ts +++ b/tfjs-backend-webgl/src/gpgpu_math.ts @@ -325,7 +325,7 @@ export function runProgram( } gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i); - }; + } const outShapeLoc = binary.outShapeLocation; if (outShapeLoc) { From b98764f4663def3a9976fd1c9a1c4b5db16c7c71 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Tue, 21 Feb 2023 15:42:56 -0800 Subject: [PATCH 7/8] Fix texture manager types --- tfjs-backend-webgl/src/texture_manager.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tfjs-backend-webgl/src/texture_manager.ts b/tfjs-backend-webgl/src/texture_manager.ts index 8173aa3caad..0d6ce717add 100644 --- a/tfjs-backend-webgl/src/texture_manager.ts +++ b/tfjs-backend-webgl/src/texture_manager.ts @@ -27,8 +27,8 @@ export class TextureManager { private _numBytesAllocated = 0; // Number of bytes that have been allocated and available for reuse. private _numBytesFree = 0; - private freeTextures?: Record = {}; - private usedTextures?: Record = {}; + private freeTextures: Record = {}; + private usedTextures: Record = {}; private logEnabled = false; constructor(private readonly gpgpu: GPGPUContext) {} @@ -118,7 +118,7 @@ export class TextureManager { const texList = this.usedTextures[shapeKey]; const texIndex = texList && texList.indexOf(texture); - if (texIndex == null || texIndex < 0) { + if (texIndex < 0) { throw new Error( 'Cannot release a texture that was never provided by this ' + 'texture manager'); @@ -173,6 +173,7 @@ export class TextureManager { this.gpgpu.deleteMatrixTexture(tex.texture); }); } + // TODO: Assign non-null value (empty object) to textures after disposed. this.freeTextures = null; this.usedTextures = null; this.numUsedTextures = 0; From 7e0760ec3020cbdad30db75ba4df2404ed2e0498 Mon Sep 17 00:00:00 2001 From: Chunnien Chan Date: Tue, 21 Feb 2023 15:44:45 -0800 Subject: [PATCH 8/8] fix --- tfjs-backend-webgl/src/texture_manager.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgl/src/texture_manager.ts b/tfjs-backend-webgl/src/texture_manager.ts index 0d6ce717add..1d5ea9e44c2 100644 --- a/tfjs-backend-webgl/src/texture_manager.ts +++ b/tfjs-backend-webgl/src/texture_manager.ts @@ -118,7 +118,7 @@ export class TextureManager { const texList = this.usedTextures[shapeKey]; const texIndex = texList && texList.indexOf(texture); - if (texIndex < 0) { + if (texIndex == null || texIndex < 0) { throw new Error( 'Cannot release a texture that was never provided by this ' + 'texture manager');