Improve dl.time (#715)

`dl.time()` now returns an object with the following properties (all times in ms):
   - `wallMs`: wall execution time.
   - `kernelMs`: kernel execution time, ignoring data transfer.
   - On `WebGL` the following additional properties exist:
     - `uploadWaitMs`: cpu blocking time on texture uploads.
     - `downloadWaitMs`: cpu blocking time on texture downloads (readPixels).

Author: dsmilkov

src/engine.ts

@@ -17,7 +17,7 @@
 
 import {ENV} from './environment';
 import {tidy} from './globals';
-import {KernelBackend} from './kernels/backend';
+import {BackendTimingInfo, KernelBackend} from './kernels/backend';
 import * as kernel_registry from './kernels/kernel_registry';
 import {KernelConfigRegistry} from './kernels/kernel_registry';
 import * as ops from './ops/ops';
@@ -54,10 +54,12 @@ export interface TensorManager {
 }
 
 export type MemoryInfo = {
-  numTensors: number; numDataBuffers: number; numBytes: number; backendInfo: {};
+  numTensors: number; numDataBuffers: number; numBytes: number;
   unreliable?: boolean;
 };
 
+export interface TimingInfo extends BackendTimingInfo { wallMs: number; }
+
 export class Engine implements TensorManager {
   // Public since optimizers will use it.
   registeredVariables: NamedVariableMap = {};
@@ -168,17 +170,11 @@ export class Engine implements TensorManager {
   }
 
   memory(): MemoryInfo {
-    const backendInfo = this.backend.memory();
-    const memInfo: MemoryInfo = {
-      numTensors: this.numTensors,
-      numDataBuffers: this.numDataBuffers,
-      numBytes: this.numBytes,
-      backendInfo,
-    };
-    if (backendInfo.unreliable) {
-      memInfo.unreliable = true;
-    }
-    return memInfo;
+    const info = this.backend.memory() as MemoryInfo;
+    info.numTensors = this.numTensors;
+    info.numDataBuffers = this.numDataBuffers;
+    info.numBytes = this.numBytes;
+    return info;
   }
 
   private shouldRecord(): boolean {
@@ -358,8 +354,11 @@ export class Engine implements TensorManager {
       numChannels: number): Tensor3D {
     return this.backend.fromPixels(pixels, numChannels);
   }
-  time(query: () => void): Promise<number> {
-    return this.backend.time(query);
+  async time(query: () => void): Promise<TimingInfo> {
+    const start = performance.now();
+    const timingInfo = await this.backend.time(query) as TimingInfo;
+    timingInfo.wallMs = performance.now() - start;
+    return timingInfo;
   }
 
   /**

src/environment.ts

@@ -254,7 +254,6 @@ export class Environment {
    *    - On CPU, true. Due to automatic garbage collection, these numbers
    *     represent undisposed tensors, i.e. not wrapped in `dl.tidy()`, or
    *     lacking a call to `tensor.dispose()`.
-   * - `backendInfo`: Backend-specific information.
    */
   @doc({heading: 'Performance', subheading: 'Memory'})
   static memory(): MemoryInfo {

src/index.ts

@@ -41,7 +41,8 @@ export {ConstantInitializer, Initializer, OnesInitializer, RandomNormalInitializ
 export {AdamaxOptimizer} from './graph/optimizers/adamax_optimizer';
 export {CostReduction, FeedEntry, Session} from './graph/session';
 export {MathBackendCPU, NDArrayMathCPU} from './kernels/backend_cpu';
-export {MathBackendWebGL, NDArrayMathGPU} from './kernels/backend_webgl';
+// tslint:disable-next-line:max-line-length
+export {MathBackendWebGL, NDArrayMathGPU, WebGLTimingInfo} from './kernels/backend_webgl';
 export {MatrixOrientation} from './kernels/types/matmul';
 export {GPGPUContext} from './kernels/webgl/gpgpu_context';
 export {NDArrayMath} from './math';

src/kernels/backend.ts

@@ -21,6 +21,9 @@ import {Conv2DInfo} from '../ops/conv_util';
 import {DataId, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor';
 import {DataType, Rank, TypedArray} from '../types';
 
+// Required information for all backends.
+export interface BackendTimingInfo { kernelMs: number; }
+
 export interface TensorStorage {
   read(dataId: DataId): Promise<TypedArray>;
   readSync(dataId: DataId): TypedArray;
@@ -29,12 +32,13 @@ export interface TensorStorage {
   fromPixels(
       pixels: ImageData|HTMLImageElement|HTMLCanvasElement|HTMLVideoElement,
       numChannels: number): Tensor3D;
-  time(query: () => void): Promise<number>;
   register(dataId: DataId, shape: number[], dtype: DataType): void;
   memory(): {unreliable: boolean;};  // Backend-specific information.
 }
 
-export interface BackendTimer { time(f: () => void): Promise<number>; }
+export interface BackendTimer {
+  time(f: () => void): Promise<BackendTimingInfo>;
+}
 
 /**
  * The interface that defines the kernels that should be implemented when

src/kernels/backend_cpu.ts

@@ -30,7 +30,7 @@ import {DataId, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor'
 import * as types from '../types';
 import {DataType, DataTypeMap, Rank, TypedArray} from '../types';
 import * as util from '../util';
-import {KernelBackend} from './backend';
+import {BackendTimingInfo, KernelBackend} from './backend';
 
 export class MathBackendCPU implements KernelBackend {
   private data = new WeakMap<DataId, DataTypeMap[DataType]>();
@@ -117,10 +117,11 @@ export class MathBackendCPU implements KernelBackend {
     }
   }
 
-  async time(f: () => void): Promise<number> {
+  async time(f: () => void): Promise<BackendTimingInfo> {
     const start = performance.now();
     f();
-    return performance.now() - start;
+    const kernelMs = performance.now() - start;
+    return {kernelMs};
   }
   memory() {
     return {

src/kernels/backend_webgl.ts

@@ -15,6 +15,7 @@
  * =============================================================================
  */
 
+import {TimingInfo} from '../engine';
 import {ENV} from '../environment';
 import {NDArrayMath} from '../math';
 import * as axis_util from '../ops/axis_util';
@@ -26,6 +27,7 @@ import * as types from '../types';
 // tslint:disable-next-line:max-line-length
 import {DataType, DataTypeMap, Rank, RecursiveArray, TypedArray} from '../types';
 import * as util from '../util';
+
 import {KernelBackend} from './backend';
 import {ArgMinMaxProgram} from './webgl/argminmax_gpu';
 import {AvgPool2DBackpropProgram} from './webgl/avg_pool_backprop_gpu';
@@ -70,12 +72,21 @@ export interface CPUTimerQuery {
   endMs?: number;
 }
 
+export interface WebGLTimingInfo extends TimingInfo {
+  uploadWaitMs: number;
+  downloadWaitMs: number;
+}
+
 export class MathBackendWebGL implements KernelBackend {
   private texData = new WeakMap<DataId, TextureData>();
   private canvas: HTMLCanvasElement;
 
   private programTimersStack: TimerNode[];
   private activeTimers: TimerNode[];
+  // Accumulated time spent (including blocking) in uploading data to webgl.
+  private uploadWaitMs = 0;
+  // Accumulated time spent (including blocking in downloading data from webgl.
+  private downloadWaitMs = 0;
 
   register(dataId: DataId, shape: number[], dtype: DataType): void {
     if (this.texData.has(dataId)) {
@@ -152,8 +163,16 @@ export class MathBackendWebGL implements KernelBackend {
       this.cacheOnCPU(dataId);
       return values;
     }
+    const shouldTimeProgram = this.activeTimers != null;
+    let start: number;
+    if (shouldTimeProgram) {
+      start = performance.now();
+    }
     const float32Values =
         this.gpgpu.downloadMatrixFromTexture(texture, texShape[0], texShape[1]);
+    if (shouldTimeProgram) {
+      this.downloadWaitMs += performance.now() - start;
+    }
     this.cacheOnCPU(dataId, float32Values);
     return texData.values;
   }
@@ -182,7 +201,7 @@ export class MathBackendWebGL implements KernelBackend {
     return this.readSync(dataId);
   }
 
-  async time(f: () => void): Promise<number> {
+  async time(f: () => void): Promise<WebGLTimingInfo> {
     const oldActiveTimers = this.activeTimers;
     const newActiveTimers: TimerNode[] = [];
 
@@ -204,11 +223,20 @@ export class MathBackendWebGL implements KernelBackend {
       this.programTimersStack = null;
     }
 
-    return Promise.all(flattenedActiveTimers).then(results => {
+    const kernelMs = await Promise.all(flattenedActiveTimers).then(results => {
       let sum = 0;
       results.forEach(result => sum += result);
       return sum;
     });
+    const res: WebGLTimingInfo = {
+      uploadWaitMs: this.uploadWaitMs,
+      downloadWaitMs: this.downloadWaitMs,
+      kernelMs,
+      wallMs: null  // will be filled by the engine
+    };
+    this.uploadWaitMs = 0;
+    this.downloadWaitMs = 0;
+    return res;
   }
   memory() {
     return {unreliable: false};
@@ -933,6 +961,11 @@ export class MathBackendWebGL implements KernelBackend {
       // Array is already on GPU. No-op.
       return;
     }
+    const shouldTimeProgram = this.activeTimers != null;
+    let start: number;
+    if (shouldTimeProgram) {
+      start = performance.now();
+    }
     const texShape =
         webgl_util.getTextureShapeFromLogicalShape(this.gpgpu.gl, shape);
     texData.texShape = texShape;
@@ -945,6 +978,9 @@ export class MathBackendWebGL implements KernelBackend {
           texShape[1], typedArrayToFloat32(values, dtype));
       // Once uploaded, don't store the values on cpu.
       texData.values = null;
+      if (shouldTimeProgram) {
+        this.uploadWaitMs += performance.now() - start;
+      }
     }
   }
 

src/profiler.ts

@@ -15,12 +15,11 @@
  * =============================================================================
  */
 
-import * as util from './util';
-import {Tensor} from './tensor';
-import {TypedArray} from './types';
-
 import {BackendTimer} from './kernels/backend';
 import {Kernel} from './kernels/kernel_registry';
+import {Tensor} from './tensor';
+import {TypedArray} from './types';
+import * as util from './util';
 
 export class Profiler {
   constructor(private backendTimer: BackendTimer, private logger?: Logger) {
@@ -39,8 +38,8 @@ export class Profiler {
     const vals = result.dataSync();
     util.checkForNaN(vals, result.dtype, kernelName);
 
-    timer.then((timeMs: number) => {
-      this.logger.logKernelProfile(kernelName, result, vals, timeMs);
+    timer.then(timing => {
+      this.logger.logKernelProfile(kernelName, result, vals, timing.kernelMs);
     });
 
     return result as T;

src/profiler_test.ts

@@ -16,21 +16,22 @@
  */
 
 import * as dl from './index';
-import {Tensor} from './tensor';
-import {BackendTimer} from './kernels/backend';
+import {BackendTimer, BackendTimingInfo} from './kernels/backend';
 import {Kernel} from './kernels/kernel_registry';
-import {Logger, Profiler} from './profiler';
 import {TypedArray} from './kernels/webgl/tex_util';
+import {Logger, Profiler} from './profiler';
+import {Tensor} from './tensor';
 
 class TestBackendTimer implements BackendTimer {
   private counter = 1;
   constructor(private delayMs: number, private queryTimeMs: number) {}
 
-  time(query: () => void): Promise<number> {
+  async time(query: () => void): Promise<BackendTimingInfo> {
     query();
-    return new Promise<number>(
+    const kernelMs = await new Promise<number>(
         resolve => setTimeout(
             resolve(this.queryTimeMs * this.counter++), this.delayMs));
+    return {kernelMs};
   }
 }
 

src/tracking.ts

@@ -16,6 +16,7 @@
  */
 
 import {doc} from './doc';
+import {TimingInfo} from './engine';
 import {ENV} from './environment';
 // tslint:disable-next-line:max-line-length
 import {ScopeFn, ScopeResult, ScopeResultImmediate} from './tape_util';
@@ -88,12 +89,20 @@ export class Tracking {
   }
 
   /**
-   * Executes `f()` and returns a promise that resolves with the elapsed time of
-   * `f()` in milliseconds.
+   * Executes `f()` and returns a promise that resolves with timing information.
+   *
+   * The result is an object with the following properties:
+   *
+   * - `wallMs`: wall execution time.
+   * - `kernelMs`: kernel execution time, ignoring data transfer.
+   * - On `WebGL` the following additional properties exist:
+   *   - `uploadWaitMs`: cpu blocking time on texture uploads.
+   *   - `downloadWaitMs`: cpu blocking time on texture downloads (readPixels).
+   *
    * @param f The function to execute and time.
    */
   @doc({heading: 'Performance', subheading: 'Timing'})
-  static time(f: () => void): Promise<number> {
+  static time(f: () => void): Promise<TimingInfo> {
     return ENV.engine.time(f);
   }
 }

src/tracking_test.ts

@@ -0,0 +1,71 @@
+/**
+ * @license
+ * Copyright 2018 Google Inc. 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
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as dl from './index';
+import {CPU_ENVS, describeWithFlags, WEBGL_ENVS} from './test_util';
+
+describeWithFlags('time webgl', WEBGL_ENVS, () => {
+  it('upload + compute', async () => {
+    const a = dl.zeros([10, 10]);
+    const time = await dl.time(() => a.square()) as dl.WebGLTimingInfo;
+    expect(time.uploadWaitMs > 0);
+    expect(time.downloadWaitMs === 0);
+    expect(time.kernelMs > 0);
+    expect(time.wallMs >= time.kernelMs);
+  });
+
+  it('upload + compute + dataSync', async () => {
+    const a = dl.zeros([10, 10]);
+    const time =
+        await dl.time(() => a.square().dataSync()) as dl.WebGLTimingInfo;
+    expect(time.uploadWaitMs > 0);
+    expect(time.downloadWaitMs > 0);
+    expect(time.kernelMs > 0);
+    expect(time.wallMs >= time.kernelMs);
+  });
+
+  it('upload + compute + data', async () => {
+    const a = dl.zeros([10, 10]);
+    const time = await dl.time(async () => await a.square().data()) as
+        dl.WebGLTimingInfo;
+    expect(time.uploadWaitMs > 0);
+    expect(time.downloadWaitMs > 0);
+    expect(time.kernelMs > 0);
+    expect(time.wallMs >= time.kernelMs);
+  });
+
+  it('preupload (not included) + compute + data', async () => {
+    const a = dl.zeros([10, 10]);
+    // Pre-upload a on gpu.
+    a.square();
+    const time = await dl.time(() => a.sqrt()) as dl.WebGLTimingInfo;
+    // The tensor was already on gpu.
+    expect(time.uploadWaitMs === 0);
+    expect(time.downloadWaitMs === 0);
+    expect(time.kernelMs > 0);
+    expect(time.wallMs >= time.kernelMs);
+  });
+});
+
+describeWithFlags('time cpu', CPU_ENVS, () => {
+  it('simple upload', async () => {
+    const a = dl.zeros([10, 10]);
+    const time = await dl.time(() => a.square());
+    expect(time.kernelMs > 0);
+    expect(time.wallMs >= time.kernelMs);
+  });
+});