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 | #ifndef TENSORFLOW_DTENSOR_CC_DSTATUS_H_ |
17 | #define TENSORFLOW_DTENSOR_CC_DSTATUS_H_ |
18 | |
19 | #include <vector> |
20 | |
21 | #include "absl/strings/match.h" |
22 | #include "absl/strings/str_cat.h" |
23 | #include "absl/strings/str_format.h" |
24 | #include "absl/strings/str_join.h" |
25 | #include "absl/types/optional.h" |
26 | #include "tensorflow/compiler/xla/stream_executor/lib/statusor.h" |
27 | #include "tensorflow/core/platform/errors.h" |
28 | #include "tensorflow/core/platform/status.h" |
29 | |
30 | namespace tensorflow { |
31 | namespace dtensor { |
32 | |
33 | template <typename T> |
34 | using StatusOr = stream_executor::port::StatusOr<T>; |
35 | |
36 | inline Status WithContext(const Status& ds, absl::string_view file, |
37 | int line_number, absl::string_view context = "" ) { |
38 | if (ds.ok()) { |
39 | return ds; |
40 | } |
41 | return Status(ds.code(), absl::StrCat(ds.error_message(), "\n" , file, ":" , |
42 | line_number, " :: " , context)); |
43 | } |
44 | |
45 | template <class T> |
46 | inline StatusOr<T> WithContext(StatusOr<T>&& ds, absl::string_view file, |
47 | int line_number, |
48 | absl::string_view context = "" ) { |
49 | if (ds.ok()) { |
50 | return ds; |
51 | } |
52 | return Status(ds.status().code(), |
53 | absl::StrCat(ds.status().error_message(), "\n" , file, ":" , |
54 | line_number, " :: " , context)); |
55 | } |
56 | |
57 | #define DT_CTX(dstatus, ...) \ |
58 | ::tensorflow::dtensor::WithContext(dstatus, __FILE__, __LINE__, #__VA_ARGS__); |
59 | |
60 | #undef TF_RETURN_IF_ERROR |
61 | #define TF_RETURN_IF_ERROR(...) \ |
62 | do { \ |
63 | ::tensorflow::Status _status = (__VA_ARGS__); \ |
64 | if (!_status.ok()) { \ |
65 | return ::tensorflow::dtensor::WithContext(_status, __FILE__, __LINE__); \ |
66 | } \ |
67 | } while (0); |
68 | |
69 | #undef TF_RETURN_WITH_CONTEXT |
70 | #define TF_RETURN_WITH_CONTEXT(status, ...) \ |
71 | do { \ |
72 | ::tensorflow::Status _status = (status); \ |
73 | if (!_status.ok()) { \ |
74 | return ::tensorflow::dtensor::WithContext(_status, __FILE__, __LINE__, \ |
75 | ##__VA_ARGS__); \ |
76 | } \ |
77 | } while (0); |
78 | |
79 | #define DT_STATUS_MACROS_CONCAT_NAME(x, y) DT_STATUS_MACROS_CONCAT_IMPL(x, y) |
80 | #define DT_STATUS_MACROS_CONCAT_IMPL(x, y) x##y |
81 | |
82 | #define DT_ASSIGN_OR_RETURN_IMPL(statusor, lhs, rexpr, ...) \ |
83 | auto statusor = (rexpr); \ |
84 | if (!statusor.ok()) { \ |
85 | return ::tensorflow::dtensor::WithContext(statusor.status(), __FILE__, \ |
86 | __LINE__, ##__VA_ARGS__); \ |
87 | } \ |
88 | lhs = std::move(statusor.value()) |
89 | |
90 | #undef TF_ASSIGN_OR_RETURN |
91 | #define TF_ASSIGN_OR_RETURN(lhs, rexpr, ...) \ |
92 | DT_ASSIGN_OR_RETURN_IMPL( \ |
93 | DT_STATUS_MACROS_CONCAT_NAME(_status_or_value, __COUNTER__), lhs, rexpr, \ |
94 | ##__VA_ARGS__) |
95 | |
96 | // Undefine TF status macros to ensure users use the context macros instead |
97 | |
98 | } // namespace dtensor |
99 | } // namespace tensorflow |
100 | |
101 | #endif // TENSORFLOW_DTENSOR_CC_DSTATUS_H_ |
102 | |