1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17"""Graph debug results dumping class."""
18import collections
19import json
20import os
21import numpy as np
22import tvm
23
24GRAPH_DUMP_FILE_NAME = '_tvmdbg_graph_dump.json'
25CHROME_TRACE_FILE_NAME = "_tvmdbg_execution_trace.json"
26
27ChromeTraceEvent = collections.namedtuple(
28    'ChromeTraceEvent',
29    ['ts', 'tid', 'pid', 'name', 'ph']
30)
31
32
33class DebugResult(object):
34    """Graph debug data module.
35
36    Data dump module manage all the debug data formatting.
37    Output data and input graphs are formatted and dumped to file.
38    Frontend read these data and graph for visualization.
39
40    Parameters
41    ----------
42    graph_json : str
43        The graph to be deployed in json format output by nnvm graph. Each operator (tvm_op)
44        in the graph will have a one to one mapping with the symbol in libmod which is used
45        to construct a "PackedFunc" .
46
47    dump_path : str
48        Output data path is read/provided from frontend
49    """
50
51    def __init__(self, graph_json, dump_path):
52        self._dump_path = dump_path
53        self._output_tensor_list = []
54        self._time_list = []
55        self._parse_graph(graph_json)
56        # dump the json information
57        self.dump_graph_json(graph_json)
58
59    def _parse_graph(self, graph_json):
60        """Parse and extract the NNVM graph and update the nodes, shapes and dltype.
61
62        Parameters
63        ----------
64        graph_json : str or graph class
65           The graph to be deployed in json format output by nnvm graph.
66        """
67        json_obj = json.loads(graph_json)
68        self._nodes_list = json_obj['nodes']
69        self._shapes_list = json_obj['attrs']['shape']
70        self._dtype_list = json_obj['attrs']['dltype']
71        self._update_graph_json()
72
73    def _update_graph_json(self):
74        """update the nodes_list with name, shape and data type,
75        for temporarily storing the output.
76        """
77
78        nodes_len = len(self._nodes_list)
79        for i in range(nodes_len):
80            node = self._nodes_list[i]
81            input_list = []
82            for input_node in node['inputs']:
83                input_list.append(self._nodes_list[input_node[0]]['name'])
84            node['inputs'] = input_list
85            dtype = str("type: " + self._dtype_list[1][i])
86            if 'attrs' not in node:
87                node['attrs'] = {}
88                node['op'] = "param"
89            else:
90                node['op'] = node['attrs']['func_name']
91            node['attrs'].update({"T": dtype})
92            node['shape'] = self._shapes_list[1][i]
93
94    def _cleanup_tensors(self):
95        """Remove the tensor dump file (graph wont be removed)
96        """
97        for filename in os.listdir(self._dump_path):
98            if os.path.isfile(filename) and not filename.endswith(".json"):
99                os.remove(filename)
100
101    def get_graph_nodes(self):
102        """Return the nodes list
103        """
104        return self._nodes_list
105
106    def get_graph_node_shapes(self):
107        """Return the nodes shapes list
108        """
109        return self._shapes_list
110
111    def get_graph_node_output_num(self, node):
112        """Return the number of outputs of a node
113        """
114        return 1 if node['op'] == 'param' else int(node['attrs']['num_outputs'])
115
116    def get_graph_node_dtypes(self):
117        """Return the nodes dtype list
118        """
119        return self._dtype_list
120
121    def get_output_tensors(self):
122        """Dump the outputs to a temporary folder, the tensors are in numpy format
123        """
124        eid = 0
125        order = 0
126        output_tensors = {}
127        for node, time in zip(self._nodes_list, self._time_list):
128            num_outputs = self.get_graph_node_output_num(node)
129            for j in range(num_outputs):
130                order += time[0]
131                key = node['name'] + "_" + str(j)
132                output_tensors[key] = self._output_tensor_list[eid]
133                eid += 1
134        return output_tensors
135
136    def dump_output_tensor(self):
137        """Dump the outputs to a temporary folder, the tensors are in numpy format
138        """
139        #cleanup existing tensors before dumping
140        self._cleanup_tensors()
141        eid = 0
142        order = 0
143        output_tensors = {}
144        for node, time in zip(self._nodes_list, self._time_list):
145            num_outputs = self.get_graph_node_output_num(node)
146            for j in range(num_outputs):
147                order += time[0]
148                key = node['name'] + "_" + str(j) + "__" + str(order)
149                output_tensors[key] = self._output_tensor_list[eid]
150                eid += 1
151
152        with open(os.path.join(self._dump_path, "output_tensors.params"), "wb") as param_f:
153            param_f.write(save_tensors(output_tensors))
154
155    def dump_chrome_trace(self):
156        """Dump the trace to the Chrome trace.json format.
157        """
158        def s_to_us(t):
159            return t * 10 ** 6
160
161        starting_times = np.zeros(len(self._time_list) + 1)
162        starting_times[1:] = np.cumsum([times[0] for times in self._time_list])
163
164        def node_to_events(node, times, starting_time):
165            return [
166                ChromeTraceEvent(
167                    ts=s_to_us(starting_time),
168                    tid=1,
169                    pid=1,
170                    ph='B',
171                    name=node['name'],
172                ),
173                ChromeTraceEvent(
174                    # Use start + duration instead of end to ensure precise timings.
175                    ts=s_to_us(times[0] + starting_time),
176                    tid=1,
177                    pid=1,
178                    ph='E',
179                    name=node['name'],
180                ),
181            ]
182        events = [
183            e for (node, times, starting_time) in zip(
184                self._nodes_list, self._time_list, starting_times)
185            for e in node_to_events(node, times, starting_time)]
186        result = dict(
187            displayTimeUnit='ns',
188            traceEvents=[e._asdict() for e in events]
189        )
190
191        with open(os.path.join(self._dump_path, CHROME_TRACE_FILE_NAME), "w") as trace_f:
192            json.dump(result, trace_f)
193
194    def dump_graph_json(self, graph):
195        """Dump json formatted graph.
196
197        Parameters
198        ----------
199        graph : json format
200            json formatted NNVM graph contain list of each node's
201            name, shape and type.
202        """
203        graph_dump_file_name = GRAPH_DUMP_FILE_NAME
204        with open(os.path.join(self._dump_path, graph_dump_file_name), 'w') as outfile:
205            json.dump(graph, outfile, indent=4, sort_keys=False)
206
207    def display_debug_result(self, sort_by_time=True):
208        """Displays the debugger result"
209        """
210        header = ["Node Name", "Ops", "Time(us)", "Time(%)", "Shape", "Inputs", "Outputs"]
211        lines = ["---------", "---", "--------", "-------", "-----", "------", "-------"]
212        eid = 0
213        data = []
214        total_time = sum(time[0] for time in self._time_list)
215        for node, time in zip(self._nodes_list, self._time_list):
216            num_outputs = self.get_graph_node_output_num(node)
217            for j in range(num_outputs):
218                op = node['op']
219                if node['op'] == 'param':
220                    eid += 1
221                    continue
222                name = node['name']
223                shape = str(self._output_tensor_list[eid].shape)
224                time_us = round(time[0] * 1000000, 3)
225                time_percent = round(((time[0] / total_time) * 100), 3)
226                inputs = str(node['attrs']['num_inputs'])
227                outputs = str(node['attrs']['num_outputs'])
228                node_data = [name, op, time_us, time_percent, shape, inputs, outputs]
229                data.append(node_data)
230                eid += 1
231
232        if sort_by_time:
233            # Sort on the basis of execution time. Prints the most expensive ops in the start.
234            data = sorted(data, key=lambda x: x[2], reverse=True)
235            # Insert a row for total time at the end.
236            rounded_total_time = round(total_time * 1000000, 3)
237            data.append(["Total_time", "-", rounded_total_time, "-", "-", "-", "-", "-"])
238
239        fmt = ""
240        for i, _ in enumerate(header):
241            max_len = len(header[i])
242            for j, _ in enumerate(data):
243                item_len = len(str(data[j][i]))
244                if item_len > max_len:
245                    max_len = item_len
246            fmt = fmt + "{:<" + str(max_len + 2) + "}"
247        print(fmt.format(*header))
248        print(fmt.format(*lines))
249        for row in data:
250            print(fmt.format(*row))
251
252def save_tensors(params):
253    """Save parameter dictionary to binary bytes.
254
255    The result binary bytes can be loaded by the
256    GraphModule with API "load_params".
257
258    Parameters
259    ----------
260    params : dict of str to NDArray
261        The parameter dictionary.
262
263    Returns
264    -------
265    param_bytes: bytearray
266        Serialized parameters.
267    """
268    _save_tensors = tvm.get_global_func("_save_param_dict")
269
270    args = []
271    for k, v in params.items():
272        args.append(k)
273        args.append(tvm.nd.array(v))
274    return _save_tensors(*args)
275