1 | /* Copyright 2022 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 | #include "tensorflow/dtensor/mlir/dtensor_location.h" |
17 | |
18 | #include <algorithm> |
19 | #include <queue> |
20 | #include <string> |
21 | |
22 | #include "llvm/ADT/STLExtras.h" |
23 | #include "llvm/ADT/SmallVector.h" |
24 | #include "llvm/Support/FormatVariadic.h" |
25 | #include "llvm/Support/raw_ostream.h" |
26 | #include "tensorflow/compiler/mlir/utils/name_utils.h" |
27 | |
28 | namespace tensorflow { |
29 | namespace dtensor { |
30 | |
31 | namespace { |
32 | std::string CreateLocalLocationString(mlir::FileLineColLoc loc) { |
33 | return llvm::formatv(">> {0}:{1}:{2}" , loc.getFilename(), loc.getLine(), |
34 | loc.getColumn()) |
35 | .str(); |
36 | } |
37 | } // namespace |
38 | |
39 | mlir::Location DTensorLocation(mlir::Location loc, llvm::StringRef file, |
40 | unsigned int line, llvm::StringRef name) { |
41 | // Strip dirname. |
42 | auto split = file.rsplit("/" ); |
43 | if (!split.second.empty()) file = split.second; |
44 | mlir::Location callee_loc = |
45 | mlir::FileLineColLoc::get(loc.getContext(), file, line, 0); |
46 | std::string new_name = GetNameFromLoc(loc); |
47 | if (!new_name.empty()) { |
48 | if (!name.empty()) { |
49 | new_name = llvm::formatv("{0}/{1}" , new_name, name).str(); |
50 | } |
51 | callee_loc = mlir::NameLoc::get( |
52 | mlir::StringAttr::get(loc.getContext(), new_name), callee_loc); |
53 | } |
54 | return mlir::CallSiteLoc::get(/*callee=*/callee_loc, /*caller=*/loc); |
55 | } |
56 | |
57 | mlir::Location DTensorLocation(mlir::Operation* op, llvm::StringRef file, |
58 | unsigned int line, llvm::StringRef name) { |
59 | return DTensorLocation(op->getLoc(), file, line, name); |
60 | } |
61 | |
62 | std::string DTensorLocationToString(mlir::Location loc) { |
63 | llvm::SmallVector<std::string, 4> stack; |
64 | std::queue<mlir::Location> queue; |
65 | queue.push(loc); |
66 | |
67 | while (!queue.empty()) { |
68 | mlir::Location& front = queue.front(); |
69 | if (auto name_loc = front.dyn_cast<mlir::NameLoc>()) { |
70 | queue.push(name_loc.getChildLoc()); |
71 | } else if (auto callsite_loc = front.dyn_cast<mlir::CallSiteLoc>()) { |
72 | queue.push(callsite_loc.getCallee()); |
73 | queue.push(callsite_loc.getCaller()); |
74 | } else if (auto line_loc = front.dyn_cast<mlir::FileLineColLoc>()) { |
75 | stack.push_back(CreateLocalLocationString(line_loc)); |
76 | } |
77 | queue.pop(); |
78 | } |
79 | |
80 | std::reverse(stack.begin(), stack.end()); |
81 | std::string s; |
82 | llvm::raw_string_ostream ss(s); |
83 | llvm::interleave(stack, ss, "\n" ); |
84 | return ss.str(); |
85 | } |
86 | |
87 | } // namespace dtensor |
88 | } // namespace tensorflow |
89 | |