WebGL profiling tool (#285)

* Adding WebGL Query for profiling.

* change webgl query time from nanoseconds to miniseconds.

* adding docs

* Accomodating PR comments

docs/development.md

@@ -56,3 +56,11 @@ To debug on an Android device, please refer to the following steps:
 - Follow instructions in [Remote Debugging on Android](https://developer.chrome.com/devtools/docs/remote-debugging-legacy) to launch `adb`. Make sure to use port 9000 so that the existing debug configuration works.
 - in VSCode, select debug configuration `Remote Browser via Webkit Adaptor`.
 - follow the steps above to debug.
+
+### Performance profiling
+
+To profile the kernel level execution time, please run the following command:
+
+npm test -- model <model_name> --profile
+
+It generates raw perf data for each kernel, you may want use tools/parse-profiler.ts to parse it.

lib/backends/webgl/utils.ts

@@ -12,6 +12,31 @@ export function getPackedShape(unpackedShape: ReadonlyArray<number>): ReadonlyAr
   return unpackedShape.slice(0, len - 1).concat(unpackedShape[len - 1] / 4);
 }
 
+export function repeatedTry(
+    checkFn: () => boolean, delayFn = (counter: number) => 0, maxCounter?: number): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    let tryCount = 0;
+
+    const tryFn = () => {
+      if (checkFn()) {
+        resolve();
+        return;
+      }
+
+      tryCount++;
+
+      const nextBackoff = delayFn(tryCount);
+
+      if (maxCounter != null && tryCount >= maxCounter) {
+        reject();
+        return;
+      }
+      setTimeout(tryFn, nextBackoff);
+    };
+
+    tryFn();
+  });
+}
 /**
  * Generates the function name from an input sampler name.
  * @param samplerName Name of the sampler.

lib/backends/webgl/webgl-context.ts

@@ -5,6 +5,7 @@ import {env} from '../../env';
 
 import * as DataEncoders from './texture-data-encoder';
 import {DataEncoder, Encoder} from './texture-data-encoder';
+import {repeatedTry} from './utils';
 
 /**
  * Abstraction and wrapper around WebGLRenderingContext and its operations
@@ -41,6 +42,8 @@ export class WebGLContext {
 
   // WebGL2 extensions
   colorBufferFloatExtension: {}|null;
+  // tslint:disable-next-line: no-any
+  disjointTimerQueryWebgl2Extension: any;
 
   private disposed: boolean;
   private frameBufferBound = false;
@@ -331,6 +334,7 @@ export class WebGLContext {
   private getExtensions(): void {
     if (this.version === 2) {
       this.colorBufferFloatExtension = this.gl.getExtension('EXT_color_buffer_float');
+      this.disjointTimerQueryWebgl2Extension = this.gl.getExtension('EXT_disjoint_timer_query_webgl2');
     } else {
       this.textureFloatExtension = this.gl.getExtension('OES_texture_float');
       this.textureHalfFloatExtension = this.gl.getExtension('OES_texture_half_float');
@@ -464,4 +468,64 @@ export class WebGLContext {
       }
     }
   }
+
+  beginTimer(): WebGLQuery {
+    if (this.version === 2) {
+      const gl2 = this.gl as WebGL2RenderingContext;
+      const ext = this.disjointTimerQueryWebgl2Extension;
+
+      const query = gl2.createQuery() as WebGLQuery;
+      gl2.beginQuery(ext.TIME_ELAPSED_EXT, query);
+      return query;
+    } else {
+      // TODO: add webgl 1 handling.
+      throw new Error('WebGL1 profiling currently not supported.');
+    }
+  }
+
+  endTimer() {
+    if (this.version === 2) {
+      const gl2 = this.gl as WebGL2RenderingContext;
+      const ext = this.disjointTimerQueryWebgl2Extension;
+      gl2.endQuery(ext.TIME_ELAPSED_EXT);
+      return;
+    } else {
+      // TODO: add webgl 1 handling.
+      throw new Error('WebGL1 profiling currently not supported');
+    }
+  }
+
+  isTimerResultAvailable(query: WebGLQuery): boolean {
+    let available = false, disjoint = false;
+    if (this.version === 2) {
+      const gl2 = this.gl as WebGL2RenderingContext;
+      const ext = this.disjointTimerQueryWebgl2Extension;
+
+      available = gl2.getQueryParameter(query, gl2.QUERY_RESULT_AVAILABLE);
+      disjoint = gl2.getParameter(ext.GPU_DISJOINT_EXT);
+    } else {
+      // TODO: add webgl 1 handling.
+      throw new Error('WebGL1 profiling currently not supported');
+    }
+
+    return available && !disjoint;
+  }
+
+  getTimerResult(query: WebGLQuery): number {
+    let timeElapsed = 0;
+    if (this.version === 2) {
+      const gl2 = this.gl as WebGL2RenderingContext;
+      timeElapsed = gl2.getQueryParameter(query, gl2.QUERY_RESULT);
+    } else {
+      // TODO: add webgl 1 handling.
+      throw new Error('WebGL1 profiling currently not supported');
+    }
+    // return miliseconds
+    return timeElapsed / 1000000;
+  }
+
+  async waitForQueryAndGetTime(query: WebGLQuery): Promise<number> {
+    await repeatedTry(() => this.isTimerResultAvailable(query));
+    return this.getTimerResult(query);
+  }
 }

lib/execution-plan.ts

@@ -2,6 +2,8 @@
 // Licensed under the MIT license.
 
 import {SessionHandler} from './backend';
+import {WebGLBackend} from './backends/backend-webgl';
+import {WebGLContext} from './backends/webgl/webgl-context';
 import {Graph} from './graph';
 import {Logger, Profiler} from './instrument';
 import {Operator} from './operators';
@@ -51,6 +53,12 @@ export class ExecutionPlan {
   }
 
   execute(sessionHandler: SessionHandler, modelInputs: Tensor[]): Promise<Tensor[]> {
+    const isWebGLBackend = sessionHandler.backend instanceof WebGLBackend;
+    let glCtx: WebGLContext|undefined;
+    if (isWebGLBackend) {
+      glCtx = (sessionHandler.backend as WebGLBackend).glContext;
+    }
+
     return this.profiler.event('session', 'ExecutionPlan.execute', async () => {
       // reset mediem result
       this.reset();
@@ -95,7 +103,7 @@ export class ExecutionPlan {
             `Runing op:${thisOp.node.name} (${
                 inputTensors.map((t, i) => `'${thisOp.node.inputs[i]}': ${t.type}[${t.dims.join(',')}]`).join(', ')})`);
 
-        const outputList = await this.profiler.event('node', thisOp.node.name, async () => {
+        const execNodeFn = async () => {
           const op = thisOp.op;
           if (!op.checkInputs(inputTensors)) {
             throw new Error(`invalid inputs detected; op: ${thisOp.node.name}`);
@@ -104,7 +112,10 @@ export class ExecutionPlan {
           const result = op.run(inferenceHandler, inputTensors);
 
           return result;
-        });
+        };
+
+        const outputList = isWebGLBackend ? await this.profiler.event('node', thisOp.node.name, execNodeFn, glCtx) :
+                                            await this.profiler.event('node', thisOp.node.name, execNodeFn);
 
         // check output
         if (outputList.length !== thisOp.node.outputs.length) {

lib/instrument.ts

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license.
 
+import {WebGLContext} from './backends/webgl/webgl-context';
+
 export declare namespace Logger {
   export interface SeverityTypeMap {
     verbose: 'v';
@@ -232,17 +234,28 @@ export declare namespace Profiler {
   export type EventCategory = 'session'|'node'|'op'|'backend';
 
   export interface Event {
-    end(): void;
+    end(): void|Promise<void>;
   }
 }
+// TODO
+// class WebGLEvent implements Profiler.Event {}
 
 class Event implements Profiler.Event {
   constructor(
       public category: Profiler.EventCategory, public name: string, public startTime: number,
-      private endCallback: (e: Event) => void) {}
+      private endCallback: (e: Event) => void|Promise<void>, public timer?: WebGLQuery, public ctx?: WebGLContext) {}
 
   end() {
-    this.endCallback(this);
+    return this.endCallback(this);
+  }
+
+  async checkTimer(): Promise<number> {
+    if (this.ctx === undefined || this.timer === undefined) {
+      throw new Error('No webgl timer found');
+    } else {
+      this.ctx.endTimer();
+      return this.ctx.waitForQueryAndGetTime(this.timer);
+    }
   }
 }
 
@@ -283,11 +296,12 @@ export class Profiler {
   }
 
   // create an event scope for the specific function
-  event<T>(category: Profiler.EventCategory, name: string, func: () => T): T;
-  event<T>(category: Profiler.EventCategory, name: string, func: () => Promise<T>): Promise<T>;
+  event<T>(category: Profiler.EventCategory, name: string, func: () => T, ctx?: WebGLContext): T;
+  event<T>(category: Profiler.EventCategory, name: string, func: () => Promise<T>, ctx?: WebGLContext): Promise<T>;
 
-  event<T>(category: Profiler.EventCategory, name: string, func: () => T | Promise<T>): T|Promise<T> {
-    const event = this._started ? this.begin(category, name) : undefined;
+  event<T>(category: Profiler.EventCategory, name: string, func: () => T | Promise<T>, ctx?: WebGLContext): T
+      |Promise<T> {
+    const event = this._started ? this.begin(category, name, ctx) : undefined;
     let isPromise = false;
 
     try {
@@ -299,44 +313,67 @@ export class Profiler {
         return new Promise<T>((resolve, reject) => {
           (res as Promise<T>)
               .then(
-                  value => {  // fulfilled
-                    resolve(value);
+                  async value => {  // fulfilled
                     if (event) {
-                      event.end();
+                      await event.end();
                     }
+                    resolve(value);
                   },
-                  reason => {  // rejected
-                    reject(reason);
+                  async reason => {  // rejected
                     if (event) {
-                      event.end();
+                      await event.end();
                     }
+                    reject(reason);
                   });
         });
       }
-
+      if (!isPromise && event) {
+        const eventRes = event.end();
+        if (eventRes && typeof eventRes.then === 'function') {
+          return new Promise<T>((resolve, reject) => {
+            (eventRes).then(
+                () => {  // fulfilled
+                  resolve(res);
+                },
+                (reason) => {  // rejected
+                  reject(reason);
+                });
+          });
+        }
+      }
       return res;
 
     } finally {
-      if (!isPromise && event) {
-        event.end();
-      }
     }
   }
 
   // begin an event
-  begin(category: Profiler.EventCategory, name: string): Event {
+  begin(category: Profiler.EventCategory, name: string, ctx?: WebGLContext): Event {
     if (!this._started) {
       throw new Error('profiler is not started yet');
     }
-    const startTime = now();
-    this.flush(startTime);
-    return new Event(category, name, startTime, e => this.end(e));
+    if (ctx === undefined) {
+      const startTime = now();
+      this.flush(startTime);
+      return new Event(category, name, startTime, e => this.endSync(e));
+    } else {
+      const timer: WebGLQuery = ctx.beginTimer();
+      return new Event(category, name, 0, e => this.end(e), timer, ctx);
+    }
   }
 
   // end the specific event
-  private end(event: Event) {
+  private async end(event: Event): Promise<void> {
+    const endTime: number = await event.checkTimer();
+    if (this._timingEvents.length < this._maxNumberEvents) {
+      this._timingEvents.push(new EventRecord(event.category, event.name, event.startTime, endTime));
+      this.flush(endTime);
+    }
+  }
+
+  private endSync(event: Event): void {
+    const endTime: number = now();
     if (this._timingEvents.length < this._maxNumberEvents) {
-      const endTime = now();
       this._timingEvents.push(new EventRecord(event.category, event.name, event.startTime, endTime));
       this.flush(endTime);
     }