1 /* Copyright 2019 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 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
17 #define TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
18
19 #include <string>
20 #include <vector>
21
22 #include "absl/strings/match.h"
23 #include "absl/strings/string_view.h"
24 #include "tensorflow/core/platform/macros.h"
25
26 namespace tensorflow {
27 namespace profiler {
28
29 // Special op types.
30 TF_CONST_INIT extern const absl::string_view kUnknownOp;
31 TF_CONST_INIT extern const absl::string_view kDatasetOp;
32 TF_CONST_INIT extern const absl::string_view kMemcpyHToDOp;
33 TF_CONST_INIT extern const absl::string_view kMemcpyDToHOp;
34 TF_CONST_INIT extern const absl::string_view kMemcpyDToDOp;
35 TF_CONST_INIT extern const absl::string_view kMemcpyHToHOp;
36
37 enum class Category {
38 kUnknown,
39 kTensorFlow,
40 kJax,
41 kTfData,
42 kMemcpyHToD,
43 kMemcpyDToH,
44 kMemcpyDToD,
45 kMemcpyHToH,
46 };
47
48 // Breaks a TensorFlow op fullname into name and type.
49 struct TfOp {
50 Category category = Category::kUnknown;
51 absl::string_view name;
52 absl::string_view type;
53 };
54 TfOp ParseTfOpFullname(absl::string_view tf_op_fullname);
55
56 // Returns a vector of TF name scopes extracted from a TF op name.
57 std::vector<absl::string_view> ParseTfNameScopes(absl::string_view tf_op_name);
58 std::vector<absl::string_view> ParseTfNameScopes(const TfOp& tf_op);
59
60 // Trace event name for TF ops is the op type so they have the same color in
61 // trace viewer.
62 std::string TfOpEventName(const TfOp& tf_op);
63 std::string TfOpEventName(absl::string_view tf_op_fullname);
64
65 // Trace event name for dataset ops.
66 std::string DatasetOpEventName(absl::string_view full_name);
67
68 // Returns the iterator name without prefix and parent iterator names.
69 std::string IteratorName(absl::string_view full_name);
70
71 // Returns true if the given name is a TensorFlow Dataset Op.
IsDatasetOp(absl::string_view tf_op_type)72 inline bool IsDatasetOp(absl::string_view tf_op_type) {
73 return tf_op_type == kDatasetOp;
74 }
IsDatasetOp(const TfOp & tf_op)75 inline bool IsDatasetOp(const TfOp& tf_op) {
76 return tf_op.category == Category::kTfData;
77 }
78
79 // Returns true if the given name is a TensorFlow Infeed Enqueue Op.
80 // See: tensorflow/core/tpu/kernels/infeed_ops.h
IsInfeedEnqueueOp(absl::string_view tf_op_type)81 inline bool IsInfeedEnqueueOp(absl::string_view tf_op_type) {
82 return absl::StartsWith(tf_op_type, "InfeedEnqueue");
83 }
IsInfeedEnqueueOp(const TfOp & tf_op)84 inline bool IsInfeedEnqueueOp(const TfOp& tf_op) {
85 return tf_op.category == Category::kTensorFlow &&
86 IsInfeedEnqueueOp(tf_op.type);
87 }
88
89 // Returns true if the given op is for outside compilation.
IsOutsideCompilationOp(absl::string_view tf_op_fullname,absl::string_view hlo_expression)90 inline bool IsOutsideCompilationOp(absl::string_view tf_op_fullname,
91 absl::string_view hlo_expression) {
92 if (absl::EndsWith(tf_op_fullname, ":XlaSendToHost")) return true;
93 if (absl::StrContains(hlo_expression, "send-done") &&
94 absl::StrContains(hlo_expression, "is_host_transfer=true"))
95 return true;
96 return false;
97 }
98
99 // Returns true if the given name is a TensorFlow embedding op.
IsEmbeddingOp(absl::string_view tf_op_fullname)100 inline bool IsEmbeddingOp(absl::string_view tf_op_fullname) {
101 return absl::StrContains(tf_op_fullname, "Embedding");
102 }
103
104 // Returns true if the given op is for copying data from host to device.
IsMemcpyHToDOp(absl::string_view tf_op_type)105 inline bool IsMemcpyHToDOp(absl::string_view tf_op_type) {
106 return tf_op_type == kMemcpyHToDOp;
107 }
IsMemcpyHToDOp(const TfOp & tf_op)108 inline bool IsMemcpyHToDOp(const TfOp& tf_op) {
109 return tf_op.category == Category::kMemcpyHToD;
110 }
111
112 // Returns true if the given op is for copying data from device to host.
IsMemcpyDToHOp(const TfOp & tf_op)113 inline bool IsMemcpyDToHOp(const TfOp& tf_op) {
114 return tf_op.category == Category::kMemcpyDToH;
115 }
116
117 // Returns true if the given op is for copying data from device to device.
IsMemcpyDToDOp(const TfOp & tf_op)118 inline bool IsMemcpyDToDOp(const TfOp& tf_op) {
119 return tf_op.category == Category::kMemcpyDToD;
120 }
121
122 // Returns true if the given op is for copying data from host to host.
IsMemcpyHToHOp(const TfOp & tf_op)123 inline bool IsMemcpyHToHOp(const TfOp& tf_op) {
124 return tf_op.category == Category::kMemcpyHToH;
125 }
126
127 // Splits a string of tensor shapes in "(shape1;shape2;...)" format, i.e.,
128 // delimited by '(' and ')' and separated by ';', into the individual shapes.
129 std::vector<absl::string_view> ParseTensorShapes(
130 absl::string_view tensor_shapes);
131
132 // Returns true if the given string matches OpDef.name pattern.
133 bool IsTfOpName(absl::string_view op_name);
134
135 // Returns true if the given string matches NodeDef.name pattern.
136 bool IsTfOpType(absl::string_view op_type);
137
138 // Returns true if the given string matches JAX pattern.
139 bool IsJaxOpType(absl::string_view op_type);
140
141 // Returns true if the given strings match JAX pattern.
142 bool IsJaxOpNameAndType(absl::string_view op_name, absl::string_view op_type);
143
144 } // namespace profiler
145 } // namespace tensorflow
146
147 #endif // TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
148