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)