diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index e3b8970cf5a..7804901be77 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -1276,24 +1276,20 @@ export class MathBackendWebGL extends KernelBackend { } getUniformLocations() { - for (const [, binary] of Object.entries(this.binaryCache)) { + for (const binary of Object.values(this.binaryCache)) { const { - uniformLocations, + variablesLocations, customUniformLocations, infLoc, nanLoc, - inShapesLocations, - inTexShapesLocations, outShapeLocation, outShapeStridesLocation, outTexShapeLocation } = getUniformLocations(this.gpgpu, binary.program, binary.webGLProgram); - binary.uniformLocations = uniformLocations; + binary.variablesLocations = variablesLocations; 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..9de19971ad1 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; + variablesLocations?: GPGPUVariableLocations[]; +} + +export interface GPGPUVariableLocations { + name: string; + uniform: WebGLUniformLocation; + offset: WebGLUniformLocation; + shape?: WebGLUniformLocation; + texShape?: WebGLUniformLocation; } export interface TensorData { @@ -131,12 +128,10 @@ export function compileProgram( webGLProgram, inShapeInfos, outShapeInfo, - uniformLocations: null, + variablesLocations: null, 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 variablesLocations: GPGPUVariableLocations[] = []; const customUniformLocations: WebGLUniformLocation[] = []; let outShapeLocation: WebGLUniformLocation; let outTexShapeLocation: WebGLUniformLocation; @@ -165,18 +158,21 @@ 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) { + const varLocs: GPGPUVariableLocations = { + name: varName, + uniform: gpgpu.getUniformLocation(webGLProgram, varName, shouldThrow), + offset: gpgpu.getUniformLocation( + webGLProgram, `offset${varName}`, shouldThrow), + }; if (program.enableShapeUniforms) { - inShapesLocations[`${varName}Shape`] = gpgpu.getUniformLocation( + varLocs.shape = gpgpu.getUniformLocation( webGLProgram, `${varName}Shape`, shouldThrow); - inTexShapesLocations[`${varName}TexShape`] = gpgpu.getUniformLocation( + varLocs.texShape = gpgpu.getUniformLocation( webGLProgram, `${varName}TexShape`, shouldThrow); } + + variablesLocations.push(varLocs); } if (program.enableShapeUniforms) { @@ -189,19 +185,17 @@ 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 { - uniformLocations, + variablesLocations, customUniformLocations, infLoc, nanLoc, - inShapesLocations, - inTexShapesLocations, outShapeLocation, outShapeStridesLocation, outTexShapeLocation @@ -271,12 +265,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.variablesLocations[i]; if (varShapeLoc) { const {uniformShape} = shader_compiler.getUniformInfoFromShape( @@ -298,6 +294,7 @@ export function runProgram( break; } } + if (varTexShapeLoc) { gpgpu.gl.uniform2i( varTexShapeLoc, input.texData.texShape[0], input.texData.texShape[1]); @@ -305,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) { @@ -319,7 +316,7 @@ export function runProgram( } gpgpu.gl.uniform1fv(varLoc, vals); } - return; + continue; } // If the input was sliced, upload the flat offset index. @@ -328,7 +325,7 @@ export function runProgram( } gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i); - }); + } const outShapeLoc = binary.outShapeLocation; if (outShapeLoc) { @@ -375,7 +372,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') { @@ -397,7 +395,7 @@ export function runProgram( } else { throw Error(`uniform type ${d.type} is not supported yet.`); } - }); + } } gpgpu.executeProgram(); } diff --git a/tfjs-backend-webgl/src/texture_manager.ts b/tfjs-backend-webgl/src/texture_manager.ts index 1d2d174c6a9..1d5ea9e44c2 100644 --- a/tfjs-backend-webgl/src/texture_manager.ts +++ b/tfjs-backend-webgl/src/texture_manager.ts @@ -25,13 +25,13 @@ 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: {[shape: string]: Texture[]} = {}; + // Number of bytes that have been allocated and available for reuse. + private _numBytesFree = 0; + 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; } @@ -117,13 +117,14 @@ 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'); } - texList.splice(texIndex, 1); + texList[texIndex] = texList[texList.length - 1]; + texList.pop(); this.log(); } @@ -172,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;