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 #include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
17
18 #include "absl/strings/match.h"
19 #include "absl/strings/str_cat.h"
20 #include "tensorflow/compiler/mlir/tensorflow/utils/parse_text_proto.h"
21 #include "tensorflow/core/framework/tensor.pb.h"
22 #include "tensorflow/core/framework/tensor_shape.pb.h"
23 #include "tensorflow/core/framework/types.h"
24 #include "tensorflow/core/lib/core/errors.h"
25 #include "tensorflow/core/platform/protobuf.h"
26
27 namespace tensorflow {
28 namespace mangling_util {
29 namespace {
30
31 const char kAttributePrefix[] = "tf.";
32 const char kDataTypePrefix[] = "tfdtype$";
33 const char kTensorShapePrefix[] = "tfshape$";
34 const char kTensorPrefix[] = "tftensor$";
35
36 } // namespace
37
MangleAttributeName(absl::string_view str)38 string MangleAttributeName(absl::string_view str) {
39 return absl::StrCat(kAttributePrefix, str);
40 }
41
IsMangledAttributeName(absl::string_view str)42 bool IsMangledAttributeName(absl::string_view str) {
43 return absl::StartsWith(str, kAttributePrefix);
44 }
45
DemangleAttributeName(absl::string_view str)46 absl::string_view DemangleAttributeName(absl::string_view str) {
47 DCHECK(IsMangledAttributeName(str));
48 return str.substr(std::strlen(kAttributePrefix));
49 }
50
GetMangledKind(absl::string_view str)51 MangledKind GetMangledKind(absl::string_view str) {
52 if (absl::StartsWith(str, kDataTypePrefix)) {
53 return MangledKind::kDataType;
54 } else if (absl::StartsWith(str, kTensorShapePrefix)) {
55 return MangledKind::kTensorShape;
56 } else if (absl::StartsWith(str, kTensorPrefix)) {
57 return MangledKind::kTensor;
58 } else {
59 return MangledKind::kUnknown;
60 }
61 }
62
MangleShape(const TensorShapeProto & shape)63 string MangleShape(const TensorShapeProto& shape) {
64 return absl::StrCat(kTensorShapePrefix, shape.ShortDebugString());
65 }
66
DemangleShape(absl::string_view str,TensorShapeProto * proto)67 Status DemangleShape(absl::string_view str, TensorShapeProto* proto) {
68 return ParseTextProto(str, kTensorShapePrefix, proto);
69 }
70
MangleTensor(const TensorProto & tensor)71 string MangleTensor(const TensorProto& tensor) {
72 return absl::StrCat(kTensorPrefix, tensor.ShortDebugString());
73 }
74
DemangleTensor(absl::string_view str,TensorProto * proto)75 Status DemangleTensor(absl::string_view str, TensorProto* proto) {
76 return ParseTextProto(str, kTensorPrefix, proto);
77 }
78
MangleDataType(const DataType & dtype)79 string MangleDataType(const DataType& dtype) {
80 return absl::StrCat(kDataTypePrefix, DataType_Name(dtype));
81 }
82
DemangleDataType(absl::string_view str,DataType * proto)83 Status DemangleDataType(absl::string_view str, DataType* proto) {
84 absl::string_view pbtxt;
85 TF_RETURN_IF_ERROR(ConsumePrefix(str, kDataTypePrefix, &pbtxt));
86 if (!DataType_Parse(string(pbtxt), proto)) {
87 return errors::FailedPrecondition(
88 "Could not parse TFDataType mangled proto");
89 }
90 return Status::OK();
91 }
92
93 } // namespace mangling_util
94 } // namespace tensorflow
95