Add ndarray.getValuesAsync() method which returns a promise that resolves when values are ready.

demos/benchmarks/conv_benchmarks.ts

@@ -80,8 +80,8 @@ export class ConvGPUBenchmark extends ConvBenchmark {
         gpgpu.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        gpgpu.runBenchmark(benchmark).then((timeElapsed: number) => {
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        gpgpu.runQuery(benchmark).then((timeElapsed: number) => {
           delayedCleanup();
           resolve(timeElapsed);
         });

demos/benchmarks/conv_transposed_benchmarks.ts

@@ -81,8 +81,8 @@ export class ConvTransposedGPUBenchmark extends ConvTransposedBenchmark {
         gpgpu.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        gpgpu.runBenchmark(benchmark).then((timeElapsed: number) => {
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        gpgpu.runQuery(benchmark).then((timeElapsed: number) => {
           delayedCleanup();
           resolve(timeElapsed);
         });

demos/benchmarks/logsumexp_benchmarks.ts

@@ -66,8 +66,8 @@ export class LogSumExpGPUBenchmark extends BenchmarkTest {
         gpgpu.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        gpgpu.runBenchmark(benchmark).then((timeElapsed: number) => {
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        gpgpu.runQuery(benchmark).then((timeElapsed: number) => {
           delayedCleanup();
           resolve(timeElapsed);
         });

demos/benchmarks/matmul_benchmarks.ts

@@ -80,8 +80,8 @@ export class MatmulGPUBenchmark extends BenchmarkTest {
         gpgpu.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        gpgpu.runBenchmark(benchmark).then((timeElapsed: number) => {
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        gpgpu.runQuery(benchmark).then((timeElapsed: number) => {
           delayedCleanup();
           resolve(timeElapsed);
         });

demos/benchmarks/pool_benchmarks.ts

@@ -96,8 +96,8 @@ export class PoolGPUBenchmark extends PoolBenchmark {
         gpgpu.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        gpgpu.runBenchmark(benchmark).then((timeElapsed: number) => {
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        gpgpu.runQuery(benchmark).then((timeElapsed: number) => {
           delayedCleanup();
           resolve(timeElapsed);
         });

demos/benchmarks/reduction_ops_benchmark.ts

@@ -75,8 +75,8 @@ export class ReductionOpsGPUBenchmark extends ReductionOpsBenchmark {
         math.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        math.getGPGPUContext().runBenchmark(benchmark).then(
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        math.getGPGPUContext().runQuery(benchmark).then(
             (timeElapsed: number) => {
               delayedCleanup();
               resolve(timeElapsed);

demos/benchmarks/unary_ops_benchmark.ts

@@ -103,8 +103,8 @@ export class UnaryOpsGPUBenchmark extends UnaryOpsBenchmark {
         math.dispose();
       };
 
-      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER')) {
-        math.getGPGPUContext().runBenchmark(benchmark).then(
+      if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
+        math.getGPGPUContext().runQuery(benchmark).then(
             (timeElapsed: number) => {
               delayedCleanup();
               resolve(timeElapsed);

src/environment.ts

@@ -24,12 +24,18 @@ export enum Type {
 }
 
 export interface Features {
-  'WEBGL_DISJOINT_QUERY_TIMER'?: boolean;
+  // Whether the disjoint_query_timer extension is an available extension.
+  'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED'?: boolean;
+  // Whether the timer object from the disjoint_query_timer extension gives
+  // timing information that is reliable.
+  'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE'?: boolean;
+  // 0: No WebGL, 1: WebGL 1.0, 2: WebGL 2.0.
   'WEBGL_VERSION'?: number;
 }
 
 export const URL_PROPERTIES: URLProperty[] = [
-  {name: 'WEBGL_DISJOINT_QUERY_TIMER', type: Type.BOOLEAN},
+  {name: 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED', type: Type.BOOLEAN},
+  {name: 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE', type: Type.BOOLEAN},
   {name: 'WEBGL_VERSION', type: Type.NUMBER}
 ];
 
@@ -38,50 +44,51 @@ export interface URLProperty {
   type: Type;
 }
 
-function isWebGL2Enabled() {
+function getWebGLRenderingContext(webGLVersion: number): WebGLRenderingContext {
+  if (webGLVersion === 0) {
+    throw new Error('Cannot get WebGL rendering context, WebGL is disabled.');
+  }
+
   const tempCanvas = document.createElement('canvas');
-  const gl = tempCanvas.getContext('webgl2') as WebGLRenderingContext;
+  if (webGLVersion === 1) {
+    return (tempCanvas.getContext('webgl') ||
+            tempCanvas.getContext('experimental-webgl')) as
+        WebGLRenderingContext;
+  }
+  return tempCanvas.getContext('webgl2') as WebGLRenderingContext;
+}
+
+function loseContext(gl: WebGLRenderingContext) {
   if (gl != null) {
     const loseContextExtension = gl.getExtension('WEBGL_lose_context');
     if (loseContextExtension == null) {
       throw new Error(
           'Extension WEBGL_lose_context not supported on this browser.');
     }
     loseContextExtension.loseContext();
-    return true;
   }
-  return false;
 }
 
-function isWebGL1Enabled() {
-  const tempCanvas = document.createElement('canvas');
-  const gl =
-      (tempCanvas.getContext('webgl') ||
-       tempCanvas.getContext('experimental-webgl')) as WebGLRenderingContext;
+function isWebGLVersionEnabled(webGLVersion: 1|2) {
+  const gl = getWebGLRenderingContext(webGLVersion);
   if (gl != null) {
-    const loseContextExtension = gl.getExtension('WEBGL_lose_context');
-    if (loseContextExtension == null) {
-      throw new Error(
-          'Extension WEBGL_lose_context not supported on this browser.');
-    }
-    loseContextExtension.loseContext();
+    loseContext(gl);
     return true;
   }
   return false;
 }
 
-function evaluateFeature<K extends keyof Features>(feature: K): Features[K] {
-  if (feature === 'WEBGL_DISJOINT_QUERY_TIMER') {
-    return !device_util.isMobile();
-  } else if (feature === 'WEBGL_VERSION') {
-    if (isWebGL2Enabled()) {
-      return 2;
-    } else if (isWebGL1Enabled()) {
-      return 1;
-    }
-    return 0;
+function isWebGLDisjointQueryTimerEnabled(webGLVersion: number) {
+  const gl = getWebGLRenderingContext(webGLVersion);
+
+  const extensionName = webGLVersion === 1 ? 'EXT_disjoint_timer_query' :
+                                             'EXT_disjoint_timer_query_webgl2';
+  const ext = gl.getExtension(extensionName);
+  const isExtEnabled = ext != null;
+  if (gl != null) {
+    loseContext(gl);
   }
-  throw new Error(`Unknown feature ${feature}.`);
+  return isExtEnabled;
 }
 
 export class Environment {
@@ -98,10 +105,33 @@ export class Environment {
       return this.features[feature];
     }
 
-    this.features[feature] = evaluateFeature(feature);
+    this.features[feature] = this.evaluateFeature(feature);
 
     return this.features[feature];
   }
+
+  private evaluateFeature<K extends keyof Features>(feature: K): Features[K] {
+    if (feature === 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED') {
+      const webGLVersion = this.get('WEBGL_VERSION');
+
+      if (webGLVersion === 0) {
+        return false;
+      }
+
+      return isWebGLDisjointQueryTimerEnabled(webGLVersion);
+    } else if (feature === 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE') {
+      return this.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED') &&
+          !device_util.isMobile();
+    } else if (feature === 'WEBGL_VERSION') {
+      if (isWebGLVersionEnabled(2)) {
+        return 2;
+      } else if (isWebGLVersionEnabled(1)) {
+        return 1;
+      }
+      return 0;
+    }
+    throw new Error(`Unknown feature ${feature}.`);
+  }
 }
 
 // Expects flags from URL in the format ?dljsflags=FLAG1:1,FLAG2:true.

src/environment_test.ts

@@ -15,23 +15,101 @@
  * =============================================================================
  */
 import * as device_util from './device_util';
-import {Environment} from './environment';
+import {Environment, Features} from './environment';
 
-describe('disjoint query timer', () => {
-  it('mobile', () => {
+describe('disjoint query timer enabled', () => {
+  it('no webgl', () => {
+    const features: Features = {'WEBGL_VERSION': 0};
+
+    const env = new Environment(features);
+
+    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED')).toBe(false);
+  });
+
+  it('webgl 1', () => {
+    const features: Features = {'WEBGL_VERSION': 1};
+
+    spyOn(document, 'createElement').and.returnValue({
+      getContext: (context: string) => {
+        if (context === 'webgl' || context === 'experimental-webgl') {
+          return {
+            getExtension: (extensionName: string) => {
+              if (extensionName === 'EXT_disjoint_timer_query') {
+                return {};
+              } else if (extensionName === 'WEBGL_lose_context') {
+                return {loseContext: () => {}};
+              }
+              return null;
+            }
+          };
+        }
+        return null;
+      }
+    });
+
+    const env = new Environment(features);
+
+    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED')).toBe(true);
+  });
+
+  it('webgl 2', () => {
+    const features: Features = {'WEBGL_VERSION': 2};
+
+    spyOn(document, 'createElement').and.returnValue({
+      getContext: (context: string) => {
+        if (context === 'webgl2') {
+          return {
+            getExtension: (extensionName: string) => {
+              if (extensionName === 'EXT_disjoint_timer_query_webgl2') {
+                return {};
+              } else if (extensionName === 'WEBGL_lose_context') {
+                return {loseContext: () => {}};
+              }
+              return null;
+            }
+          };
+        }
+        return null;
+      }
+    });
+
+    const env = new Environment(features);
+
+    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED')).toBe(true);
+  });
+
+
+});
+describe('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE', () => {
+  it('disjoint query timer disabled', () => {
+    const features:
+        Features = {'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED': false};
+
+    const env = new Environment(features);
+
+    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE'))
+        .toBe(false);
+  });
+
+  it('disjoint query timer enabled, mobile', () => {
+    const features:
+        Features = {'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED': true};
     spyOn(device_util, 'isMobile').and.returnValue(true);
 
-    const env = new Environment();
+    const env = new Environment(features);
 
-    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER')).toBe(false);
+    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE'))
+        .toBe(false);
   });
 
-  it('not mobile', () => {
+  it('disjoint query timer enabled, not mobile', () => {
+    const features:
+        Features = {'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED': true};
     spyOn(device_util, 'isMobile').and.returnValue(false);
 
-    const env = new Environment();
+    const env = new Environment(features);
 
-    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER')).toBe(true);
+    expect(env.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')).toBe(true);
   });
 });
 

src/math/ndarray.ts

@@ -15,6 +15,7 @@
  * =============================================================================
  */
 
+import {ENV} from '../environment';
 import * as util from '../util';
 
 import {GPGPUContext} from './webgl/gpgpu_context';
@@ -243,6 +244,27 @@ export class NDArray {
     return this.data.values;
   }
 
+  getValuesAsync(): Promise<Float32Array> {
+    return new Promise<Float32Array>((resolve, reject) => {
+      if (this.data.values != null) {
+        resolve(this.data.values);
+        return;
+      }
+
+      if (!ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED')) {
+        resolve(this.getValues());
+        return;
+      }
+
+      // Construct an empty query. We're just interested in getting a callback
+      // when the GPU command queue has executed until this point in time.
+      const queryFn = () => {};
+      GPGPU.runQuery(queryFn).then(() => {
+        resolve(this.getValues());
+      });
+    });
+  }
+
   private uploadToGPU(preferredTexShape?: [number, number]) {
     throwIfGPUNotInitialized();
     this.data.textureShapeRC = webgl_util.getTextureShapeFromLogicalShape(