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/core/kernels/tensor_cord.h"
17
18 #include <cstring>
19
20 #include "tensorflow/core/framework/variant.h"
21
22 namespace tensorflow {
23
24 static_assert(Variant::CanInlineType<TensorCord>(),
25 "TensorCord should be inlined into Variants");
26
~CordRep()27 TensorCord::CordRep::~CordRep() {
28 if (!is_inline_ && rep_.external.releaser) {
29 rep_.external.releaser(rep_.external.arg);
30 }
31 }
32
~TensorCord()33 TensorCord::~TensorCord() { Cleanup(); }
34
Encode(VariantTensorData * data) const35 void TensorCord::Encode(VariantTensorData* data) const {
36 data->metadata_string().clear();
37 for (auto rep : Chunks()) {
38 data->metadata_string().append(rep.data(), rep.size());
39 }
40 }
41
Decode(VariantTensorData data)42 bool TensorCord::Decode(VariantTensorData data) {
43 auto* str = new string(std::move(data.metadata_string()));
44 Cleanup();
45 chunks_.push_back(new CordRep(absl::string_view(*str), &StringReleaser, str));
46 return true;
47 }
48
TensorBufWithRef(Tensor * tensor)49 TensorBuffer* TensorCord::TensorBufWithRef(Tensor* tensor) {
50 TensorBuffer* buf = tensor->buf_;
51 buf->Ref();
52 return buf;
53 }
54
TensorBufReleaser(void * tensor_buffer)55 void TensorCord::TensorBufReleaser(void* tensor_buffer) {
56 static_cast<TensorBuffer*>(tensor_buffer)->Unref();
57 }
58
StringReleaser(void * str_ptr)59 void TensorCord::StringReleaser(void* str_ptr) {
60 delete static_cast<string*>(str_ptr);
61 }
62
63 namespace {
64
65 // Helpers for STLStringResizeUninitialized
66 // HasMember is true_type or false_type, depending on whether or not
67 // T has a __resize_default_init member. Resize will call the
68 // __resize_default_init member if it exists, and will call the resize
69 // member otherwise.
70 template <typename string_type, typename = void>
71 struct ResizeUninitializedTraits {
72 using HasMember = std::false_type;
Resizetensorflow::__anon9689349c0111::ResizeUninitializedTraits73 static void Resize(string_type* s, size_t new_size) { s->resize(new_size); }
74 };
75
76 // __resize_default_init is provided by libc++ >= 8.0.
77 template <typename string_type>
78 struct ResizeUninitializedTraits<
79 string_type, absl::void_t<decltype(std::declval<string_type&>()
80 .__resize_default_init(237))> > {
81 using HasMember = std::true_type;
Resizetensorflow::__anon9689349c0111::ResizeUninitializedTraits82 static void Resize(string_type* s, size_t new_size) {
83 s->__resize_default_init(new_size);
84 }
85 };
86
87 // Resize string `s` to `new_size`, leaving the data uninitialized.
STLStringResizeUninitialized(string * s,size_t new_size)88 static inline void STLStringResizeUninitialized(string* s, size_t new_size) {
89 ResizeUninitializedTraits<string>::Resize(s, new_size);
90 }
91
92 } // namespace
93
operator string() const94 TensorCord::operator string() const {
95 string out;
96 STLStringResizeUninitialized(&out, size());
97 char* data = const_cast<char*>(out.data());
98 for (auto* rep : chunks_) {
99 auto view = rep->view();
100 memcpy(data, view.data(), view.size());
101 data += view.size();
102 }
103 DCHECK_EQ(data - out.data(), size());
104 return out;
105 }
106
107 } // namespace tensorflow
108