[DO-NOT-REVIEW] webgl auto-version/feature WIP

src/backends/webgl/backend_webgl.ts

@@ -62,7 +62,7 @@ import * as binaryop_gpu from './binaryop_gpu';
 import {BinaryOpProgram} from './binaryop_gpu';
 import * as binaryop_packed_gpu from './binaryop_packed_gpu';
 import {BinaryOpPackedProgram} from './binaryop_packed_gpu';
-import {createCanvas, getWebGLContext} from './canvas_util';
+import {createCanvas} from './canvas_util';
 import {ClipProgram} from './clip_gpu';
 import {ClipPackedProgram} from './clip_packed_gpu';
 import {ComplexAbsProgram} from './complex_abs_gpu';
@@ -130,6 +130,7 @@ import {UnaryOpProgram} from './unaryop_gpu';
 import * as unary_packed_op from './unaryop_packed_gpu';
 import {UnaryOpPackedProgram} from './unaryop_packed_gpu';
 import {UnpackProgram} from './unpack_gpu';
+import {getActiveContext} from './webgl_context_manager';
 import * as webgl_util from './webgl_util';
 
 type KernelInfo = {
@@ -221,7 +222,6 @@ export class MathBackendWebGL implements KernelBackend {
   private dataRefCount = new WeakMap<DataId, number>();
   private numBytesInGPU = 0;
 
-  private canvas: HTMLCanvasElement;
   private fromPixels2DContext: CanvasRenderingContext2D|
       OffscreenCanvasRenderingContext2D;
 
@@ -248,15 +248,12 @@ export class MathBackendWebGL implements KernelBackend {
     }
 
     if (gpgpu == null) {
-      const gl = getWebGLContext(ENV.getNumber('WEBGL_VERSION'));
       this.binaryCache = getBinaryCache(ENV.getNumber('WEBGL_VERSION'));
-      this.gpgpu = new GPGPUContext(gl);
-      this.canvas = gl.canvas;
+      this.gpgpu = new GPGPUContext();
       this.gpgpuCreatedLocally = true;
     } else {
       this.binaryCache = {};
       this.gpgpuCreatedLocally = false;
-      this.canvas = gpgpu.gl.canvas;
     }
     this.textureManager = new TextureManager(this.gpgpu);
     this.numMBBeforeWarning = numMBBeforeWarning();
@@ -2500,11 +2497,7 @@ export class MathBackendWebGL implements KernelBackend {
       return;
     }
     this.textureManager.dispose();
-    if (this.canvas != null && this.canvas.remove != null) {
-      this.canvas.remove();
-    } else {
-      this.canvas = null;
-    }
+
     if (this.fromPixels2DContext != null &&
         //@ts-ignore
         this.fromPixels2DContext.canvas.remove) {
@@ -2515,6 +2508,11 @@ export class MathBackendWebGL implements KernelBackend {
       this.gpgpu.program = null;
       this.gpgpu.dispose();
     }
+
+    const gl = getActiveContext();
+    if (gl.canvas != null && gl.canvas.remove != null) {
+      gl.canvas.remove();
+    }
     this.disposed = true;
   }
 

src/backends/webgl/canvas_util.ts

@@ -1,6 +1,6 @@
 /**
  * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
+ * Copyright 2019 Google LLC. All Rights Reserved.
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
@@ -15,8 +15,6 @@
  * =============================================================================
  */
 
-const contexts: {[key: string]: WebGLRenderingContext} = {};
-
 const WEBGL_ATTRIBUTES: WebGLContextAttributes = {
   alpha: false,
   antialias: false,
@@ -27,34 +25,6 @@ const WEBGL_ATTRIBUTES: WebGLContextAttributes = {
   failIfMajorPerformanceCaveat: true
 };
 
-export function setWebGLContext(
-    webGLVersion: number, gl: WebGLRenderingContext) {
-  contexts[webGLVersion] = gl;
-}
-
-export function getWebGLContext(webGLVersion: number): WebGLRenderingContext {
-  if (!(webGLVersion in contexts)) {
-    contexts[webGLVersion] = getWebGLRenderingContext(webGLVersion);
-  }
-  const gl = contexts[webGLVersion];
-  if (gl.isContextLost()) {
-    delete contexts[webGLVersion];
-    return getWebGLContext(webGLVersion);
-  }
-
-  gl.disable(gl.DEPTH_TEST);
-  gl.disable(gl.STENCIL_TEST);
-  gl.disable(gl.BLEND);
-  gl.disable(gl.DITHER);
-  gl.disable(gl.POLYGON_OFFSET_FILL);
-  gl.disable(gl.SAMPLE_COVERAGE);
-  gl.enable(gl.SCISSOR_TEST);
-  gl.enable(gl.CULL_FACE);
-  gl.cullFace(gl.BACK);
-
-  return contexts[webGLVersion];
-}
-
 export function createCanvas(webGLVersion: number) {
   if (typeof OffscreenCanvas !== 'undefined' && webGLVersion === 2) {
     return new OffscreenCanvas(300, 150);
@@ -65,15 +35,26 @@ export function createCanvas(webGLVersion: number) {
   }
 }
 
-function getWebGLRenderingContext(webGLVersion: number): WebGLRenderingContext {
+export function cleanupDOMCanvasWebGLRenderingContext(
+    context: WebGLRenderingContext) {
+  if (context == null) {
+    throw new Error('Shold not hit this case');
+  }
+  const canvas = context.canvas;
+  if (canvas != null && canvas.remove != null) {
+    canvas.remove();
+  }
+}
+
+export function createDOMCanvasWebGLRenderingContext(webGLVersion: number):
+    WebGLRenderingContext {
   if (webGLVersion !== 1 && webGLVersion !== 2) {
     throw new Error('Cannot get WebGL rendering context, WebGL is disabled.');
   }
   const canvas = createCanvas(webGLVersion);
 
   canvas.addEventListener('webglcontextlost', (ev: Event) => {
     ev.preventDefault();
-    delete contexts[webGLVersion];
   }, false);
   if (webGLVersion === 1) {
     return (canvas.getContext('webgl', WEBGL_ATTRIBUTES) ||

src/backends/webgl/canvas_util_test.ts

@@ -1,6 +1,6 @@
 /**
  * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
+ * Copyright 2019 Google LLC. All Rights Reserved.
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
@@ -18,27 +18,29 @@
 import {ENV} from '../../environment';
 import {BROWSER_ENVS, describeWithFlags} from '../../jasmine_util';
 
-import {getWebGLContext} from './canvas_util';
+import {createDOMCanvasWebGLRenderingContext} from './canvas_util';
 
 describeWithFlags('canvas_util', BROWSER_ENVS, () => {
   it('Returns a valid canvas', () => {
-    const canvas = getWebGLContext(ENV.getNumber('WEBGL_VERSION')).canvas as (
-                       HTMLCanvasElement | OffscreenCanvas);
+    const canvas =
+        createDOMCanvasWebGLRenderingContext(ENV.getNumber('WEBGL_VERSION'))
+            .canvas as (HTMLCanvasElement | OffscreenCanvas);
     expect(
         (canvas instanceof HTMLCanvasElement) ||
         (canvas instanceof OffscreenCanvas))
         .toBe(true);
   });
 
   it('Returns a valid gl context', () => {
-    const gl = getWebGLContext(ENV.getNumber('WEBGL_VERSION'));
+    const gl =
+        createDOMCanvasWebGLRenderingContext(ENV.getNumber('WEBGL_VERSION'));
     expect(gl.isContextLost()).toBe(false);
   });
 });
 
 describeWithFlags('canvas_util webgl2', {flags: {WEBGL_VERSION: 2}}, () => {
   it('is ok when the user requests webgl 1 canvas', () => {
-    const canvas = getWebGLContext(1).canvas;
+    const canvas = createDOMCanvasWebGLRenderingContext(1).canvas;
     expect((canvas instanceof HTMLCanvasElement)).toBe(true);
   });
 });

src/backends/webgl/clip_gpu.ts

@@ -17,6 +17,7 @@
 
 import {GPGPUContext} from './gpgpu_context';
 import {GPGPUProgram} from './gpgpu_math';
+import {getActiveContext} from './webgl_context_manager';
 
 export class ClipProgram implements GPGPUProgram {
   variableNames = ['A'];
@@ -51,8 +52,9 @@ export class ClipProgram implements GPGPUProgram {
         this.minLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'min');
         this.maxLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'max');
       }
-      gpgpu.gl.uniform1f(this.minLoc, min);
-      gpgpu.gl.uniform1f(this.maxLoc, max);
+      const gl = getActiveContext();
+      gl.uniform1f(this.minLoc, min);
+      gl.uniform1f(this.maxLoc, max);
     };
   }
 }

src/backends/webgl/clip_packed_gpu.ts

@@ -17,6 +17,7 @@
 
 import {GPGPUContext} from './gpgpu_context';
 import {GPGPUProgram} from './gpgpu_math';
+import {getActiveContext} from './webgl_context_manager';
 
 export class ClipPackedProgram implements GPGPUProgram {
   variableNames = ['A'];
@@ -53,8 +54,9 @@ export class ClipPackedProgram implements GPGPUProgram {
         this.minLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'min');
         this.maxLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'max');
       }
-      gpgpu.gl.uniform1f(this.minLoc, min);
-      gpgpu.gl.uniform1f(this.maxLoc, max);
+      const gl = getActiveContext();
+      gl.uniform1f(this.minLoc, min);
+      gl.uniform1f(this.maxLoc, max);
     };
   }
 }

src/backends/webgl/fill_gpu.ts

@@ -17,6 +17,7 @@
 
 import {GPGPUContext} from './gpgpu_context';
 import {GPGPUProgram} from './gpgpu_math';
+import {getActiveContext} from './webgl_context_manager';
 
 export class FillProgram implements GPGPUProgram {
   variableNames: string[];
@@ -43,7 +44,7 @@ export class FillProgram implements GPGPUProgram {
       if (this.valueLoc == null) {
         this.valueLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'value');
       }
-      gpgpu.gl.uniform1f(this.valueLoc, value);
+      getActiveContext().uniform1f(this.valueLoc, value);
     };
   }
 }

src/backends/webgl/flags_webgl_test.ts

@@ -19,7 +19,7 @@ import * as device_util from '../../device_util';
 import {ENV} from '../../environment';
 import {webgl_util} from '../../webgl';
 
-import * as canvas_util from './canvas_util';
+import * as webgl_context_manager from './webgl_context_manager';
 
 describe('HAS_WEBGL', () => {
   beforeEach(() => ENV.reset());
@@ -197,7 +197,7 @@ describe('WEBGL_MAX_TEXTURE_SIZE', () => {
     ENV.reset();
     webgl_util.MAX_TEXTURE_SIZE = null;
 
-    spyOn(canvas_util, 'getWebGLContext').and.returnValue({
+    spyOn(webgl_context_manager, 'getContextByVersion').and.returnValue({
       MAX_TEXTURE_SIZE: 101,
       getParameter: (param: number) => {
         if (param === 101) {
@@ -223,7 +223,7 @@ describe('WEBGL_MAX_TEXTURES_IN_SHADER', () => {
     ENV.reset();
     webgl_util.MAX_TEXTURES_IN_SHADER = null;
 
-    spyOn(canvas_util, 'getWebGLContext').and.callFake(() => {
+    spyOn(webgl_context_manager, 'getContextByVersion').and.callFake(() => {
       return {
         MAX_TEXTURE_IMAGE_UNITS: 101,
         getParameter: (param: number) => {