1 /* Copyright 2016 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 
16 // Node classes used for different views. They are wrappers with "show"
17 // methods.
18 //
19 // ScopeNode is for scope view. GraphNode is for graph view, CodeNode
20 // is for code view and OpNode for op view.
21 // ScopeNode and GraphNode each maps to one TFGraphNode.
22 // CodeNode and OpNode each maps to one TFMultiGraphNode.
23 
24 #ifndef TENSORFLOW_CORE_PROFILER_INTERNAL_TFPROF_NODE_SHOW_H_
25 #define TENSORFLOW_CORE_PROFILER_INTERNAL_TFPROF_NODE_SHOW_H_
26 
27 #include <algorithm>
28 #include <string>
29 #include <vector>
30 
31 #include "tensorflow/core/framework/graph.pb.h"
32 #include "tensorflow/core/lib/core/errors.h"
33 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
34 #include "tensorflow/core/profiler/internal/tfprof_node.h"
35 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
36 #include "tensorflow/core/profiler/tfprof_options.h"
37 #include "tensorflow/core/profiler/tfprof_output.pb.h"
38 
39 namespace tensorflow {
40 namespace tfprof {
41 
42 class ShowNode {
43  public:
44   explicit ShowNode(const TFGraphNode* node);
~ShowNode()45   virtual ~ShowNode() {}
46 
name()47   const string& name() const { return node->name(); }
48   GraphNodeProto* mutable_proto();
49   const GraphNodeProto& proto() const;
50 
51   void ReInit(int64 step);
52 
53   void AggregateTotalStats(ShowNode* node);
54 
55   void AddSelfToTotalStats();
56 
57   void ResetTotalStats();
58 
59   const TFGraphNode* node;
60   bool account;
61   string formatted_str;
62 
63  protected:
64   GraphNodeProto proto_;
65 };
66 
67 class GraphNode : public ShowNode {
68  public:
GraphNode(TFGraphNode * node)69   explicit GraphNode(TFGraphNode* node) : ShowNode(node) {}
70 
Trackable(int64 step)71   bool Trackable(int64 step) const { return node->trackable(step); }
72 
73   std::vector<GraphNode*> children;
74   std::vector<GraphNode*> show_children;
75 };
76 
77 class ScopeNode : public ShowNode {
78  public:
ScopeNode(const TFGraphNode * node)79   explicit ScopeNode(const TFGraphNode* node) : ShowNode(node) {}
~ScopeNode()80   ~ScopeNode() override {}
81 
82   std::vector<ScopeNode*> children;
83   std::vector<ScopeNode*> show_children;
84 };
85 
86 class ShowMultiNode {
87  public:
88   explicit ShowMultiNode(TFMultiGraphNode* node);
~ShowMultiNode()89   virtual ~ShowMultiNode() {}
90 
91   bool ReInit(int64 step, const std::vector<string>& type_regexes);
92 
name()93   const string& name() const { return node->name(); }
94   MultiGraphNodeProto* mutable_proto();
95   const MultiGraphNodeProto& proto() const;
96 
97   void AggregateTotalStats(ShowMultiNode* node);
98 
99   void AddSelfToTotalStats();
100 
101   void ResetTotalStats();
102 
103   TFMultiGraphNode* node;
104   bool account;
105   bool show;
106   string formatted_str;
107 
108  protected:
109   MultiGraphNodeProto proto_;
110 };
111 
112 class CodeNode : public ShowMultiNode {
113  public:
CodeNode(TFMultiGraphNode * node,const CallStack::Trace * trace,const string & suffix)114   CodeNode(TFMultiGraphNode* node, const CallStack::Trace* trace,
115            const string& suffix)
116       : ShowMultiNode(node), trace_(trace), suffix_(suffix) {}
~CodeNode()117   ~CodeNode() override {}
118 
AddChildren(const string & name,const CallStack::Trace * trace,const string suffix)119   CodeNode* AddChildren(const string& name, const CallStack::Trace* trace,
120                         const string suffix) {
121     auto it = children_.find(name);
122     if (it != children_.end()) {
123       return it->second.get();
124     }
125 
126     graph_children_.push_back(
127         std::unique_ptr<TFMultiGraphNode>(new TFMultiGraphNode(name)));
128     auto child = &children_[name];
129     child->reset(new CodeNode(graph_children_.back().get(), trace, suffix));
130     children.push_back(child->get());
131     return child->get();
132   }
133 
has_trace()134   bool has_trace() const { return trace_ != nullptr; }
lineno()135   const int32 lineno() const { return trace_->lineno(); }
file()136   string file() const { return trace_->file(); }
function()137   string function() const { return trace_->function() + suffix_; }
func_start_line()138   int32 func_start_line() const { return trace_->func_start_line(); }
139 
140   std::vector<CodeNode*> children;
141   std::vector<CodeNode*> show_children;
142 
143  private:
144   const CallStack::Trace* trace_;
145   string suffix_;
146   std::vector<std::unique_ptr<TFMultiGraphNode>> graph_children_;
147   std::map<string, std::unique_ptr<CodeNode>> children_;
148 };
149 
150 class OpNode : public ShowMultiNode {
151  public:
OpNode(TFMultiGraphNode * node)152   explicit OpNode(TFMultiGraphNode* node) : ShowMultiNode(node) {}
~OpNode()153   ~OpNode() override {}
154 };
155 
156 }  // namespace tfprof
157 }  // namespace tensorflow
158 
159 #endif  // TENSORFLOW_CORE_PROFILER_INTERNAL_TFPROF_NODE_SHOW_H_
160