From f44d74c56e7de9b023fec2f3cc6395ab096f3010 Mon Sep 17 00:00:00 2001
From: Andrew Tulloch <andrew@tullo.ch>
Date: Thu, 28 Mar 2019 16:03:28 -0700
Subject: [PATCH] TVM debugresult dump to Chrome Tracing

---
 python/tvm/contrib/debugger/debug_result.py   | 50 ++++++++++++++++++-
 python/tvm/contrib/debugger/debug_runtime.py  |  4 +-
 .../unittest/test_runtime_graph_debug.py      | 16 ++++++
 3 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/python/tvm/contrib/debugger/debug_result.py b/python/tvm/contrib/debugger/debug_result.py
index 101af6887c47..7e3e97e2cae1 100644
--- a/python/tvm/contrib/debugger/debug_result.py
+++ b/python/tvm/contrib/debugger/debug_result.py
@@ -1,9 +1,18 @@
 """Graph debug results dumping class."""
-import os
+import collections
 import json
+import os
+import numpy as np
 import tvm
 
 GRAPH_DUMP_FILE_NAME = '_tvmdbg_graph_dump.json'
+CHROME_TRACE_FILE_NAME = "_tvmdbg_execution_trace.json"
+
+ChromeTraceEvent = collections.namedtuple(
+    'ChromeTraceEvent',
+    ['ts', 'tid', 'pid', 'name', 'ph']
+)
+
 
 class DebugResult(object):
     """Graph debug data module.
@@ -127,6 +136,45 @@ def dump_output_tensor(self):
         with open(os.path.join(self._dump_path, "output_tensors.params"), "wb") as param_f:
             param_f.write(save_tensors(output_tensors))
 
+    def dump_chrome_trace(self):
+        """Dump the trace to the Chrome trace.json format.
+        """
+        def s_to_us(t):
+            return t * 10 ** 6
+
+        starting_times = np.zeros(len(self._time_list) + 1)
+        starting_times[1:] = np.cumsum([times[0] for times in self._time_list])
+
+        def node_to_events(node, times, starting_time):
+            return [
+                ChromeTraceEvent(
+                    ts=s_to_us(starting_time),
+                    tid=1,
+                    pid=1,
+                    ph='B',
+                    name=node['name'],
+                ),
+                ChromeTraceEvent(
+                    # Use start + duration instead of end to ensure precise timings.
+                    ts=s_to_us(times[0] + starting_time),
+                    tid=1,
+                    pid=1,
+                    ph='E',
+                    name=node['name'],
+                ),
+            ]
+        events = [
+            e for (node, times, starting_time) in zip(
+                self._nodes_list, self._time_list, starting_times)
+            for e in node_to_events(node, times, starting_time)]
+        result = dict(
+            displayTimeUnit='ns',
+            traceEvents=[e._asdict() for e in events]
+        )
+
+        with open(os.path.join(self._dump_path, CHROME_TRACE_FILE_NAME), "w") as trace_f:
+            json.dump(result, trace_f)
+
     def dump_graph_json(self, graph):
         """Dump json formatted graph.
 
diff --git a/python/tvm/contrib/debugger/debug_runtime.py b/python/tvm/contrib/debugger/debug_runtime.py
index 725f212fce00..5fcabcc6507e 100644
--- a/python/tvm/contrib/debugger/debug_runtime.py
+++ b/python/tvm/contrib/debugger/debug_runtime.py
@@ -220,7 +220,9 @@ def run(self, **input_dict):
         self._run_debug()
         # Step 2. Dump the output tensors to the dump folder
         self.debug_datum.dump_output_tensor()
-        # Step 3. Display the collected information
+        # Step 3. Dump the Chrome trace to the dump folder
+        self.debug_datum.dump_chrome_trace()
+        # Step 4. Display the collected information
         self.debug_datum.display_debug_result()
 
     def run_individual(self, number, repeat=1, min_repeat_ms=0):
diff --git a/tests/python/unittest/test_runtime_graph_debug.py b/tests/python/unittest/test_runtime_graph_debug.py
index 4bbe6509c40c..4c59cb53eadd 100644
--- a/tests/python/unittest/test_runtime_graph_debug.py
+++ b/tests/python/unittest/test_runtime_graph_debug.py
@@ -64,6 +64,22 @@ def check_verify():
         #Verify the tensors are dumped
         assert(len(os.listdir(directory)) > 1)
 
+        CHROME_TRACE_FILE_NAME = '_tvmdbg_execution_trace.json'
+        assert(os.path.exists(os.path.join(directory, CHROME_TRACE_FILE_NAME)))
+
+        with open(os.path.join(directory, CHROME_TRACE_FILE_NAME)) as f:
+            trace = json.load(f)
+        assert trace["displayTimeUnit"] == "ns"
+        events = trace["traceEvents"]
+        assert len(events) == 4
+        assert all(event["ph"] in ('B', 'E') for event in events)
+        assert all(event["pid"] == 1 for event in events)
+        assert all(event["tid"] == 1 for event in events)
+        assert all(event["name"] == 'x' for event in events[:2])
+        assert all(event["name"] == 'add' for event in events[2:])
+        assert events[0]["ts"] == 0
+        assert events[0]["ph"] == 'B'
+
         #verify the output is correct
         out = mod.get_output(0, tvm.nd.empty((n,)))
         np.testing.assert_equal(out.asnumpy(), a + 1)