[core] Dispose all webgl programs when disposing the backend.

tfjs-core/src/backends/webgl/backend_webgl.ts

@@ -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,15 @@ export class MathBackendWebGL extends KernelBackend {
     if (this.disposed) {
       return;
     }
+    // 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);
+        delete this.binaryCache[key];
+      });
+    }
     this.textureManager.dispose();
     if (this.canvas != null &&
         (typeof (HTMLCanvasElement) !== 'undefined' &&

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';