Decouple logging from profiler.profileKernel

tfjs-core/src/engine.ts

@@ -19,7 +19,7 @@ import {BackendTimingInfo, DataMover, KernelBackend} from './backends/backend';
 import {Environment, setEnvironmentGlobal} from './environment';
 import {getGlobalNamespace} from './global_util';
 import {getGradient, getKernel, getKernelsForBackend, GradFunc, NamedAttrMap, TensorInfo} from './kernel_registry';
-import {Profiler} from './profiler';
+import {KernelProfile, Profiler} from './profiler';
 import {backpropagateGradients, getFilteredNodesXToY, TapeNode} from './tape';
 import {DataId, setTensorTracker, Tensor, TensorTracker, Variable} from './tensor';
 import {GradSaveFunc, NamedTensorMap, NamedVariableMap, TensorContainer} from './tensor_types';
@@ -51,7 +51,7 @@ export type MemoryInfo = {
   unreliable?: boolean; reasons: string[];
 };
 
-type KernelProfile = {
+type KernelInfo = {
   name: string; bytesAdded: number; totalBytesSnapshot: number;
   tensorsAdded: number;
   totalTensorsSnapshot: number;
@@ -61,7 +61,7 @@ type KernelProfile = {
 
 export type ProfileInfo = {
   newBytes: number; newTensors: number; peakBytes: number;
-  kernels: KernelProfile[];
+  kernels: KernelInfo[];
   result: TensorContainer;
 };
 
@@ -618,8 +618,11 @@ export class Engine implements TensorTracker, DataMover {
           if (!this.ENV.getBool('DEBUG')) {
             outputs = kernelFunc();
           } else {
-            outputs = this.profiler.profileKernel(
+            let kernelProfile: KernelProfile;
+            kernelProfile = this.profiler.profileKernel(
                 kernelName, inputs, () => kernelFunc());
+            this.profiler.logKernelProfile(kernelProfile);
+            outputs = kernelProfile.outputs;
           }
         });
 

tfjs-core/src/profiler.ts

@@ -21,6 +21,14 @@ import {NamedTensorMap} from './tensor_types';
 import {DataType, DataTypeMap, TypedArray} from './types';
 import * as util from './util';
 
+export type KernelProfile = {
+  kernelName: string,
+  outputs: Tensor[],
+  inputs: NamedTensorMap,
+  timeMs: Promise<number|{error: string}>,
+  extraInfo: Promise<string>
+};
+
 export class Profiler {
   constructor(private backendTimer: BackendTimer, private logger?: Logger) {
     if (logger == null) {
@@ -29,32 +37,44 @@ export class Profiler {
   }
 
   profileKernel(kernelName: string, inputs: NamedTensorMap, f: () => Tensor[]):
-      Tensor[] {
+      KernelProfile {
     let outputs: Tensor[];
     const holdResultWrapperFn = () => {
       outputs = f();
     };
     const timer = this.backendTimer.time(holdResultWrapperFn);
 
-    outputs.forEach(r => {
+    outputs.map(r => {
       // Dangling promise here because we don't want to propagate up
       // asynchronicity.
-      r.data().then(vals => {
-        checkComputationForErrors(vals, r.dtype, kernelName);
+      r.data().then(tensorVals => {
+        checkComputationForErrors(tensorVals, r.dtype, kernelName);
+      });
+    });
 
-        timer.then(timing => {
-          let extraInfo = '';
-          if (timing.getExtraProfileInfo != null) {
-            extraInfo = timing.getExtraProfileInfo();
-          }
+    const kernelProfile = {
+      kernelName,
+      outputs,
+      inputs,
+      timeMs: timer.then(timing => timing.kernelMs),
+      extraInfo: timer.then(
+          timing => timing.getExtraProfileInfo != null ?
+              timing.getExtraProfileInfo() :
+              '')
+    };
+    return kernelProfile;
+  }
+
+  logKernelProfile(kernelProfile: KernelProfile): void {
+    const {kernelName, outputs, timeMs, inputs, extraInfo} = kernelProfile;
 
-          this.logger.logKernelProfile(
-              kernelName, r, vals, timing.kernelMs, inputs, extraInfo);
-        });
+    outputs.forEach(result => {
+      Promise.all([result.data(), timeMs, extraInfo]).then(valueContainer => {
+        this.logger.logKernelProfile(
+            kernelName, result, valueContainer[0], valueContainer[1], inputs,
+            valueContainer[2]);
       });
     });
-
-    return outputs;
   }
 }
 

tfjs-core/src/profiler_test.ts

@@ -18,8 +18,9 @@
 import {BackendTimer, BackendTimingInfo} from './backends/backend';
 import * as tf from './index';
 import {describeWithFlags, SYNC_BACKEND_ENVS} from './jasmine_util';
-import {checkComputationForErrors, Logger, Profiler} from './profiler';
+import {checkComputationForErrors, KernelProfile, Logger, Profiler} from './profiler';
 import {Tensor} from './tensor';
+import {NamedTensorMap} from './tensor_types';
 import {TypedArray} from './types';
 
 class TestBackendTimer implements BackendTimer {
@@ -42,6 +43,32 @@ class TestLogger extends Logger {
       name: string, result: Tensor, vals: TypedArray, timeMs: number) {}
 }
 
+function promiseCheckWrapper(acturalValPromise: Promise<{}>, truthVal: {}) {
+  return acturalValPromise.then(acturalVal => {
+    expect(acturalVal).toEqual(truthVal);
+  });
+}
+
+function checkKernelProfile(acturalVal: KernelProfile, truthVal: {
+  kernelName: string,
+  outputs: Tensor[],
+  timeMs: number|{error: string},
+  inputs: NamedTensorMap,
+  extraInfo: string
+}) {
+  expect(acturalVal.kernelName).toBe(truthVal.kernelName);
+  expect(acturalVal.inputs).toBe(truthVal.inputs);
+  acturalVal.outputs.forEach((output, index) => {
+    expect(output).toBe(truthVal.outputs[index]);
+  });
+
+  const promiseContainer = [
+    promiseCheckWrapper(acturalVal.timeMs, truthVal.timeMs),
+    promiseCheckWrapper(acturalVal.extraInfo, truthVal.extraInfo),
+  ];
+  return Promise.all(promiseContainer);
+}
+
 describeWithFlags('profiler.Profiler', SYNC_BACKEND_ENVS, () => {
   it('profiles simple function', doneFn => {
     const delayMs = 5;
@@ -56,29 +83,26 @@ describeWithFlags('profiler.Profiler', SYNC_BACKEND_ENVS, () => {
     spyOn(logger, 'logKernelProfile').and.callThrough();
 
     const timeSpy = timer.time as jasmine.Spy;
-    const logKernelProfileSpy = logger.logKernelProfile as jasmine.Spy;
 
     let kernelCalled = false;
     const result = 1;
     const resultScalar = tf.scalar(result);
 
-    profiler.profileKernel('MatMul', inputs, () => {
+    const kernelProfile = profiler.profileKernel('MatMul', inputs, () => {
       kernelCalled = true;
       return [resultScalar];
     });
-
     setTimeout(() => {
       expect(timeSpy.calls.count()).toBe(1);
-
-      expect(logKernelProfileSpy.calls.count()).toBe(1);
-
-      expect(logKernelProfileSpy.calls.first().args).toEqual([
-        'MatMul', resultScalar, new Float32Array([result]), queryTimeMs, inputs,
-        extraInfo
-      ]);
-
       expect(kernelCalled).toBe(true);
-      doneFn();
+
+      checkKernelProfile(kernelProfile, {
+        kernelName: 'MatMul',
+        outputs: [resultScalar],
+        timeMs: queryTimeMs,
+        inputs,
+        extraInfo,
+      }).then(() => doneFn());
     }, delayMs * 2);
   });
 
@@ -95,37 +119,73 @@ describeWithFlags('profiler.Profiler', SYNC_BACKEND_ENVS, () => {
     spyOn(timer, 'time').and.callThrough();
     spyOn(logger, 'logKernelProfile').and.callThrough();
     const timeSpy = timer.time as jasmine.Spy;
-    const logKernelProfileSpy = logger.logKernelProfile as jasmine.Spy;
 
     let matmulKernelCalled = false;
     let maxKernelCalled = false;
     const result = 1;
     const resultScalar = tf.scalar(result);
 
-    profiler.profileKernel('MatMul', inputs, () => {
-      const result = profiler.profileKernel('Max', inputs, () => {
+    let innerKernelProfile: KernelProfile;
+    const outerKernelProfile = profiler.profileKernel('MatMul', inputs, () => {
+      innerKernelProfile = profiler.profileKernel('Max', inputs, () => {
         maxKernelCalled = true;
         return [resultScalar];
       });
       matmulKernelCalled = true;
-      return result;
+      return innerKernelProfile.outputs;
     });
 
     setTimeout(() => {
       expect(timeSpy.calls.count()).toBe(2);
+      expect(matmulKernelCalled).toBe(true);
+      expect(maxKernelCalled).toBe(true);
+
+      const checkInnerKernelProfile = checkKernelProfile(innerKernelProfile, {
+        kernelName: 'Max',
+        outputs: [resultScalar],
+        timeMs: queryTimeMs,
+        inputs,
+        extraInfo
+      });
+      const checkOuterKernelProfile = checkKernelProfile(outerKernelProfile, {
+        kernelName: 'MatMul',
+        outputs: [resultScalar],
+        timeMs: queryTimeMs * 2,
+        inputs,
+        extraInfo
+      });
+      Promise.all([checkInnerKernelProfile, checkOuterKernelProfile])
+          .then(() => doneFn());
+    }, delayMs * 2);
+  });
+
+  it('log kernelProfile', doneFn => {
+    const delayMs = 5;
+    const queryTimeMs = 10;
+    const inputs = {'x': tf.tensor1d([1])};
+    const extraInfo = '';
+    const timer = new TestBackendTimer(delayMs, queryTimeMs, extraInfo);
+    const logger = new TestLogger();
+    const profiler = new Profiler(timer, logger);
+
+    spyOn(logger, 'logKernelProfile').and.callThrough();
+    const logKernelProfileSpy = logger.logKernelProfile as jasmine.Spy;
+
+    const result = 1;
+    const resultScalar = tf.scalar(result);
+
+    const kernelProfiles = profiler.profileKernel('MatMul', inputs, () => {
+      return [resultScalar];
+    });
+    profiler.logKernelProfile(kernelProfiles);
+
+    setTimeout(() => {
+      expect(logKernelProfileSpy.calls.count()).toBe(1);
 
-      expect(logKernelProfileSpy.calls.count()).toBe(2);
       expect(logKernelProfileSpy.calls.first().args).toEqual([
-        'Max', resultScalar, new Float32Array([result]), queryTimeMs, inputs,
+        'MatMul', resultScalar, new Float32Array([result]), queryTimeMs, inputs,
         extraInfo
       ]);
-      expect(logKernelProfileSpy.calls.argsFor(1)).toEqual([
-        'MatMul', resultScalar, new Float32Array([result]), queryTimeMs * 2,
-        inputs, extraInfo
-      ]);
-
-      expect(matmulKernelCalled).toBe(true);
-      expect(maxKernelCalled).toBe(true);
       doneFn();
     }, delayMs * 2);
   });