1 | /* Copyright 2020 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_UTIL_MANAGED_STACK_TRACE_H_ |
17 | #define TENSORFLOW_CORE_UTIL_MANAGED_STACK_TRACE_H_ |
18 | |
19 | #include <functional> |
20 | #include <string> |
21 | #include <unordered_map> |
22 | #include <vector> |
23 | |
24 | #include "absl/strings/match.h" |
25 | #include "absl/types/optional.h" |
26 | #include "tensorflow/core/platform/stack_frame.h" |
27 | |
28 | namespace tensorflow { |
29 | |
30 | // Returns "true" on filenames which should be skipped. |
31 | using StackTraceFilter = std::function<bool(const char*)>; |
32 | |
33 | using SourceLoc = std::pair<std::string, int>; |
34 | |
35 | // Using absl::Hash breaks NVCC under Windows :P |
36 | struct PairHash { |
37 | template <class T1, class T2> |
38 | std::size_t operator()(const std::pair<T1, T2>& pair) const { |
39 | std::size_t h1 = std::hash<T1>()(pair.first); |
40 | std::size_t h2 = std::hash<T2>()(pair.second); |
41 | return h1 + 0x9e3779b9 + (h2 << 6) + (h2 >> 2); |
42 | } |
43 | }; |
44 | |
45 | // Maps filename/line_no combination into a stack frame. |
46 | using SourceMap = std::unordered_map<SourceLoc, StackFrame, PairHash>; |
47 | |
48 | using ToStackFramesFunctor = std::vector<StackFrame>(int, const SourceMap&, |
49 | const StackTraceFilter&, |
50 | bool, int); |
51 | |
52 | // Returns whether the given frame is internal to TF. |
53 | inline bool IsInternalFrameForFilename(absl::string_view file_name) { |
54 | // Use a simple heuristic for now. |
55 | // TODO(cheshire): Build a more sophisticated mechanism, rely on @tf.export. |
56 | return (absl::StrContains(file_name, "tensorflow/python" ) || |
57 | absl::StrContains(file_name, "tensorflow\\python" )) && |
58 | !absl::StrContains(file_name, "keras" ) && |
59 | !absl::StrContains(file_name, "test.py" ); |
60 | } |
61 | |
62 | // Language agnostic stack trace class. It only saves an id, and language |
63 | // clients are responsible for managing the actual stack trace objects. |
64 | class ManagedStackTrace { |
65 | public: |
66 | ManagedStackTrace(int id, ToStackFramesFunctor* to_stack_frames) |
67 | : id_(id), to_stack_frames_(to_stack_frames) {} |
68 | |
69 | // Returns stack trace as a vector of `StackFrame`s. |
70 | std::vector<StackFrame> ToStackFrames(const SourceMap& source_map, |
71 | const StackTraceFilter& filtered, |
72 | bool reverse_traversal = false, |
73 | int limit = -1) const { |
74 | return to_stack_frames_(id_, source_map, filtered, reverse_traversal, |
75 | limit); |
76 | } |
77 | |
78 | private: |
79 | int id_; |
80 | ToStackFramesFunctor* to_stack_frames_; |
81 | }; |
82 | |
83 | // Generates a message with a definition location based on a provided stack |
84 | // trace, or an empty one if the stack trace is empty. |
85 | inline std::string DefinitionLocationMsg( |
86 | const absl::optional<ManagedStackTrace>& stack_trace) { |
87 | if (stack_trace.has_value()) { |
88 | std::vector<StackFrame> stack_frames = |
89 | stack_trace->ToStackFrames({}, IsInternalFrameForFilename, |
90 | /*reverse_traversal=*/true, |
91 | /*limit=*/1); |
92 | if (!stack_frames.empty()) { |
93 | const StackFrame& last_frame = stack_frames[0]; |
94 | return absl::StrCat(" (defined @ " , last_frame.file_name, ":" , |
95 | last_frame.line_number, ")" ); |
96 | } |
97 | } |
98 | return "" ; |
99 | } |
100 | |
101 | } // namespace tensorflow |
102 | |
103 | #endif // TENSORFLOW_CORE_UTIL_MANAGED_STACK_TRACE_H_ |
104 | |