1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Data structures and algorithms for profiling information."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import os
22
23
24class ProfileDatum(object):
25  """Profile data point."""
26
27  def __init__(self,
28               device_name,
29               node_exec_stats,
30               file_path,
31               line_number,
32               func_name,
33               op_type):
34    """Constructor.
35
36    Args:
37      device_name: (string) name of the device.
38      node_exec_stats: `NodeExecStats` proto.
39      file_path: path to the source file involved in creating the op.
40      line_number: line number in the file involved in creating the op.
41      func_name: name of the function that the line belongs to.
42      op_type: (string) Operation type.
43    """
44    self.device_name = device_name
45    self.node_exec_stats = node_exec_stats
46    self.file_path = file_path
47    self.line_number = line_number
48    self.func_name = func_name
49    if self.file_path:
50      self.file_line_func = "%s:%d(%s)" % (
51          os.path.basename(self.file_path), self.line_number, self.func_name)
52    else:
53      self.file_line_func = ""
54    self.op_type = op_type
55    self.start_time = self.node_exec_stats.all_start_micros
56    self.op_time = (self.node_exec_stats.op_end_rel_micros -
57                    self.node_exec_stats.op_start_rel_micros)
58
59  @property
60  def exec_time(self):
61    """Op execution time plus pre- and post-processing."""
62    return self.node_exec_stats.all_end_rel_micros
63
64
65class AggregateProfile(object):
66  """Profile summary data for aggregating a number of ProfileDatum."""
67
68  def __init__(self, profile_datum):
69    """Constructor.
70
71    Args:
72      profile_datum: (`ProfileDatum`) an instance of `ProfileDatum` to
73        initialize this object with.
74    """
75
76    self.total_op_time = profile_datum.op_time
77    self.total_exec_time = profile_datum.exec_time
78    device_and_node = "%s:%s" % (profile_datum.device_name,
79                                 profile_datum.node_exec_stats.node_name)
80    self._node_to_exec_count = {device_and_node: 1}
81
82  def add(self, profile_datum):
83    """Accumulate a new instance of ProfileDatum.
84
85    Args:
86      profile_datum: (`ProfileDatum`) an instance of `ProfileDatum` to
87        accumulate to this object.
88    """
89
90    self.total_op_time += profile_datum.op_time
91    self.total_exec_time += profile_datum.exec_time
92    device_and_node = "%s:%s" % (profile_datum.device_name,
93                                 profile_datum.node_exec_stats.node_name)
94
95    device_and_node = "%s:%s" % (profile_datum.device_name,
96                                 profile_datum.node_exec_stats.node_name)
97    if device_and_node in self._node_to_exec_count:
98      self._node_to_exec_count[device_and_node] += 1
99    else:
100      self._node_to_exec_count[device_and_node] = 1
101
102  @property
103  def node_count(self):
104    return len(self._node_to_exec_count)
105
106  @property
107  def node_exec_count(self):
108    return sum(self._node_to_exec_count.values())
109