TVM debugresult dump to Chrome Tracing

Author: ajtulloch

python/tvm/contrib/debugger/debug_result.py

@@ -1,9 +1,12 @@
 """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"
 
 class DebugResult(object):
     """Graph debug data module.
@@ -127,6 +130,52 @@ 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.
+        """
+        eid = 0
+        order = 0
+
+        def s_to_us(t):
+            return t * 10 ** 6
+
+        TraceEvent = collections.namedtuple(
+            'TraceEvent',
+            ['ts', 'tid', 'pid', 'name', 'ph']
+        )
+
+        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 [
+                TraceEvent(
+                    ts=s_to_us(starting_time),
+                    tid=1,
+                    pid=1,
+                    ph='B',
+                    name=node['name'],
+                ),
+                TraceEvent(
+                    # 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.
 

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):

tests/python/unittest/test_runtime_graph_debug.py

@@ -6,6 +6,7 @@
 from tvm.contrib import util
 from tvm.contrib.debugger import debug_runtime as graph_runtime
 
+
 def test_graph_simple():
     n = 4
     A = tvm.placeholder((n,), name='A')
@@ -57,13 +58,31 @@ def check_verify():
         GRAPH_DUMP_FILE_NAME = '_tvmdbg_graph_dump.json'
         assert(len(os.listdir(directory)) == 1)
 
+        assert(len(os.listdir(directory)) == 1)
+
         #verify the file name is proper
         assert(os.path.exists(os.path.join(directory, GRAPH_DUMP_FILE_NAME)))
 
         mod.run()
         #Verify the tensors are dumped
         assert(len(os.listdir(directory)) > 1)
 
+        print(os.listdir(directory))
+        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)