From 9e8dcc619b1885e546abe7c73d044cbd4aa3e94b Mon Sep 17 00:00:00 2001 From: Daniel Smilkov Date: Tue, 17 Dec 2019 14:36:49 -0500 Subject: [PATCH 1/3] save --- tfjs-core/src/backends/webgl/backend_webgl.ts | 11 ++- .../src/backends/webgl/backend_webgl_test.ts | 70 ++++++++++++++++++- 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts index 3bec48180ce..1c329d188a8 100644 --- a/tfjs-core/src/backends/webgl/backend_webgl.ts +++ b/tfjs-core/src/backends/webgl/backend_webgl.ts @@ -20,7 +20,7 @@ import './flags_webgl'; import * as device_util from '../../device_util'; import {ENGINE, MemoryInfo, TimingInfo} from '../../engine'; -import {env} from '../../environment'; +import {env, ENV} from '../../environment'; import {tidy} from '../../globals'; import {TensorInfo} from '../../kernel_registry'; import {warn} from '../../log'; @@ -155,7 +155,7 @@ export interface WebGLTimingInfo extends TimingInfo { const binaryCaches: {[webGLVersion: string]: {[key: string]: GPGPUBinary}} = {}; -function getBinaryCache(webGLVersion: number) { +export function getBinaryCache(webGLVersion: number) { if (webGLVersion in binaryCaches) { return binaryCaches[webGLVersion]; } @@ -2629,6 +2629,13 @@ export class MathBackendWebGL extends KernelBackend { if (this.disposed) { return; } + if (!ENV.getBool('IS_TEST')) { + const allKeys = Object.keys(this.binaryCache); + allKeys.forEach(key => { + this.gpgpu.deleteProgram(this.binaryCache[key].webGLProgram); + delete this.binaryCache[key]; + }); + } this.textureManager.dispose(); if (this.canvas != null && (typeof (HTMLCanvasElement) !== 'undefined' && diff --git a/tfjs-core/src/backends/webgl/backend_webgl_test.ts b/tfjs-core/src/backends/webgl/backend_webgl_test.ts index d066f3ae68e..dae055fb7fd 100644 --- a/tfjs-core/src/backends/webgl/backend_webgl_test.ts +++ b/tfjs-core/src/backends/webgl/backend_webgl_test.ts @@ -21,7 +21,7 @@ import {describeWithFlags} from '../../jasmine_util'; import {expectArraysClose, expectArraysEqual} from '../../test_util'; import {decodeString} from '../../util'; -import {MathBackendWebGL, WebGLMemoryInfo} from './backend_webgl'; +import {getBinaryCache, MathBackendWebGL, WebGLMemoryInfo} from './backend_webgl'; import {WEBGL_ENVS} from './backend_webgl_test_registry'; function decodeStrings(bytes: Uint8Array[]): string[] { @@ -223,6 +223,74 @@ describeWithFlags('backendWebGL', WEBGL_ENVS, () => { }); }); +describeWithFlags('Webgl backend disposal', WEBGL_ENVS, () => { + it('register and dispose a backend outside unit test', () => { + // Simulate outside unit test environment. + tf.ENV.set('IS_TEST', false); + + const backend = new MathBackendWebGL(); + tf.registerBackend('test-disposal', () => backend); + tf.setBackend('test-disposal'); + // Compile and run a program. + tf.zeros([1000]).sqrt().dataSync(); + + // Dispose the backend. + tf.backend().dispose(); + + // Make sure the cache is empty. + const cache = getBinaryCache(tf.ENV.getNumber('WEBGL_VERSION')); + expect(Object.keys(cache).length).toBe(0); + tf.removeBackend('test-disposal'); + }); + + it('register and dispose a backend inside unit test', () => { + // Simulate inside unit test environment. + tf.ENV.set('IS_TEST', true); + + const backend = new MathBackendWebGL(); + tf.registerBackend('test-disposal', () => backend); + tf.setBackend('test-disposal'); + // Compile and run a program. + tf.zeros([1000]).sqrt().dataSync(); + + // Dispose the backend. + tf.backend().dispose(); + + // Make sure the cache is NOT empty. + const cache = getBinaryCache(tf.ENV.getNumber('WEBGL_VERSION')); + expect(Object.keys(cache).length).toBeGreaterThan(0); + tf.removeBackend('test-disposal'); + }); + + it('register, dispose and re-register a backend outside unit test', () => { + // Simulate outside unit test environment. + tf.ENV.set('IS_TEST', false); + + tf.registerBackend('test-disposal', () => new MathBackendWebGL()); + tf.setBackend('test-disposal'); + // Compile and run a program. + tf.zeros([1000]).sqrt().dataSync(); + + // Dispose the backend. + tf.backend().dispose(); + tf.removeBackend('test-disposal'); + + // Re-register a backend. + tf.registerBackend('test-disposal', () => new MathBackendWebGL()); + tf.setBackend('test-disposal'); + // Compile and run a program. + tf.zeros([1000]).sqrt().dataSync(); + + // Dispose the 2nd backend. + tf.backend().dispose(); + + // Make sure the cache is empty. + const cache = getBinaryCache(tf.ENV.getNumber('WEBGL_VERSION')); + expect(Object.keys(cache).length).toBe(0); + tf.removeBackend('test-disposal'); + }); +}); + describeWithFlags('Custom window size', WEBGL_ENVS, () => { const customBackendName = 'custom-webgl'; From 17bd87ffa15944ea693c32af0c3720e56b10d88c Mon Sep 17 00:00:00 2001 From: Daniel Smilkov Date: Tue, 17 Dec 2019 14:39:01 -0500 Subject: [PATCH 2/3] save --- tfjs-core/src/backends/webgl/backend_webgl.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts index 1c329d188a8..a14dd8a3cb9 100644 --- a/tfjs-core/src/backends/webgl/backend_webgl.ts +++ b/tfjs-core/src/backends/webgl/backend_webgl.ts @@ -2629,6 +2629,8 @@ export class MathBackendWebGL extends KernelBackend { if (this.disposed) { return; } + // We avoid disposing the compiled webgl programs during unit testing to + // avoid slowing down test execution. if (!ENV.getBool('IS_TEST')) { const allKeys = Object.keys(this.binaryCache); allKeys.forEach(key => { From 39f73a39bb3a66e331046c611d45a14c62300300 Mon Sep 17 00:00:00 2001 From: Daniel Smilkov Date: Tue, 17 Dec 2019 16:46:43 -0500 Subject: [PATCH 3/3] save --- tfjs-core/src/backends/webgl/backend_webgl.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts index a14dd8a3cb9..99b38ae9dd8 100644 --- a/tfjs-core/src/backends/webgl/backend_webgl.ts +++ b/tfjs-core/src/backends/webgl/backend_webgl.ts @@ -20,7 +20,7 @@ import './flags_webgl'; import * as device_util from '../../device_util'; import {ENGINE, MemoryInfo, TimingInfo} from '../../engine'; -import {env, ENV} from '../../environment'; +import {env} from '../../environment'; import {tidy} from '../../globals'; import {TensorInfo} from '../../kernel_registry'; import {warn} from '../../log'; @@ -2629,9 +2629,9 @@ export class MathBackendWebGL extends KernelBackend { if (this.disposed) { return; } - // We avoid disposing the compiled webgl programs during unit testing to - // avoid slowing down test execution. - if (!ENV.getBool('IS_TEST')) { + // Avoid disposing the compiled webgl programs during unit testing because + // it slows down test execution. + if (!env().getBool('IS_TEST')) { const allKeys = Object.keys(this.binaryCache); allKeys.forEach(key => { this.gpgpu.deleteProgram(this.binaryCache[key].webGLProgram);