• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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/c/tf_tensor.h"
17 
18 #include <memory>
19 
20 #include "tensorflow/c/tf_status.h"
21 #include "tensorflow/c/tf_status_helper.h"
22 #include "tensorflow/c/tf_tensor_internal.h"
23 #include "tensorflow/core/framework/allocation_description.pb.h"
24 #include "tensorflow/core/framework/log_memory.h"
25 #include "tensorflow/core/framework/tensor.h"
26 #include "tensorflow/core/framework/tensor_shape.pb.h"
27 #include "tensorflow/core/framework/types.pb.h"
28 #include "tensorflow/core/lib/core/coding.h"
29 
30 using tensorflow::Status;
31 using tensorflow::Tensor;
32 using tensorflow::TensorBuffer;
33 using tensorflow::errors::FailedPrecondition;
34 using tensorflow::errors::InvalidArgument;
35 
36 namespace tensorflow {
allocate_tensor(const char * operation,size_t len,Allocator * allocator)37 void* allocate_tensor(const char* operation, size_t len, Allocator* allocator) {
38   void* data = allocator->AllocateRaw(EIGEN_MAX_ALIGN_BYTES, len);
39   if (LogMemory::IsEnabled() && data != nullptr) {
40     LogMemory::RecordRawAllocation(
41         operation, LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, len, data,
42         allocator);
43   }
44   return data;
45 }
46 
allocate_tensor(const char * operation,size_t len)47 void* allocate_tensor(const char* operation, size_t len) {
48   return allocate_tensor(operation, len, cpu_allocator());
49 }
50 
deallocate_buffer(void * data,size_t len,void * arg)51 void deallocate_buffer(void* data, size_t len, void* arg) {
52   Allocator* allocator = nullptr;
53   if (arg == nullptr) {
54     allocator = cpu_allocator();
55   } else {
56     allocator = reinterpret_cast<Allocator*>(arg);
57   }
58   if (LogMemory::IsEnabled() && data != nullptr) {
59     LogMemory::RecordRawDeallocation(
60         "TensorFlow C Api", LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, data,
61         allocator, false);
62   }
63   allocator->DeallocateRaw(data);
64 }
65 }  // namespace tensorflow
66 
67 
TF_AllocateTensor(TF_DataType dtype,const int64_t * dims,int num_dims,size_t len)68 TF_Tensor* TF_AllocateTensor(TF_DataType dtype, const int64_t* dims,
69                              int num_dims, size_t len) {
70   void* data = tensorflow::allocate_tensor("TF_AllocateTensor", len,
71                                            tensorflow::cpu_allocator());
72   return TF_NewTensor(dtype, dims, num_dims, data, len,
73                       tensorflow::deallocate_buffer,
74                       tensorflow::cpu_allocator());
75 }
76 
TF_NewTensor(TF_DataType dtype,const int64_t * dims,int num_dims,void * data,size_t len,void (* deallocator)(void * data,size_t len,void * arg),void * deallocator_arg)77 TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims,
78                         void* data, size_t len,
79                         void (*deallocator)(void* data, size_t len, void* arg),
80                         void* deallocator_arg) {
81   std::vector<tensorflow::int64> dimvec(num_dims);
82   for (int i = 0; i < num_dims; ++i) {
83     dimvec[i] = static_cast<tensorflow::int64>(dims[i]);
84   }
85 
86   TF_ManagedBuffer* buf = nullptr;
87   if (dtype != TF_STRING && dtype != TF_RESOURCE &&
88       tensorflow::DataTypeCanUseMemcpy(
89           static_cast<tensorflow::DataType>(dtype)) &&
90       reinterpret_cast<intptr_t>(data) % std::max(1, EIGEN_MAX_ALIGN_BYTES) !=
91           0) {
92     // TF_STRING and TF_RESOURCE tensors have a different representation in
93     // TF_Tensor than they do in tensorflow::Tensor. So a copy here is a waste
94     // (any alignment requirements will be taken care of by TF_TensorToTensor
95     // and TF_TensorFromTensor).
96     //
97     // Other types have the same representation, so copy only if it is safe to
98     // do so.
99     buf = new TF_ManagedBuffer(tensorflow::allocate_tensor("TF_NewTensor", len),
100                                len, tensorflow::deallocate_buffer, nullptr);
101     std::memcpy(buf->data(), data, len);
102     // Free the original buffer.
103     deallocator(data, len, deallocator_arg);
104   } else {
105     buf = new TF_ManagedBuffer(data, len, deallocator, deallocator_arg);
106   }
107 
108   // TODO(gjn): Make the choice of interface a compile-time configuration.
109   tensorflow::TensorInterface ret(
110       Tensor(static_cast<tensorflow::DataType>(dtype),
111              tensorflow::TensorShape(dimvec), buf));
112   buf->Unref();
113   size_t elem_size = TF_DataTypeSize(dtype);
114   if (elem_size > 0 && len < (elem_size * ret.NumElements())) {
115     return nullptr;
116   }
117   return new TF_Tensor{std::make_unique<tensorflow::TensorInterface>(ret)};
118 }
119 
TF_TensorMaybeMove(TF_Tensor * t)120 TF_Tensor* TF_TensorMaybeMove(TF_Tensor* t) {
121   return t->tensor->CanMove() ? t : nullptr;
122 }
123 
TF_DeleteTensor(TF_Tensor * t)124 void TF_DeleteTensor(TF_Tensor* t) { delete t; }
125 
TF_TensorType(const TF_Tensor * t)126 TF_DataType TF_TensorType(const TF_Tensor* t) { return t->tensor->Type(); }
127 
TF_NumDims(const TF_Tensor * t)128 int TF_NumDims(const TF_Tensor* t) { return t->tensor->NumDims(); }
129 
TF_Dim(const TF_Tensor * t,int dim_index)130 int64_t TF_Dim(const TF_Tensor* t, int dim_index) {
131   return t->tensor->Dim(dim_index);
132 }
133 
TF_TensorByteSize(const TF_Tensor * t)134 size_t TF_TensorByteSize(const TF_Tensor* t) { return t->tensor->ByteSize(); }
135 
TF_TensorData(const TF_Tensor * t)136 void* TF_TensorData(const TF_Tensor* t) { return t->tensor->Data(); }
137 
TF_TensorElementCount(const TF_Tensor * t)138 int64_t TF_TensorElementCount(const TF_Tensor* t) {
139   int64_t result = 1;
140   int rank = TF_NumDims(t);
141   for (int dim = 0; dim < rank; ++dim) {
142     result *= TF_Dim(t, dim);
143   }
144   return result;
145 }
146 
TF_TensorBitcastFrom(const TF_Tensor * from,TF_DataType type,TF_Tensor * to,const int64_t * new_dims,int num_new_dims,TF_Status * status)147 void TF_TensorBitcastFrom(const TF_Tensor* from, TF_DataType type,
148                           TF_Tensor* to, const int64_t* new_dims,
149                           int num_new_dims, TF_Status* status) {
150   TF_SetStatus(status, TF_OK, "");
151   Status cc_status(
152       static_cast<tensorflow::TensorInterface*>(to->tensor.get())
153           ->BitcastFrom(*static_cast<const tensorflow::TensorInterface*>(
154                             from->tensor.get()),
155                         type, new_dims, num_new_dims));
156   Set_TF_Status_from_Status(status, cc_status);
157 }
158 
159 namespace tensorflow {
160 
CanMove() const161 bool TensorInterface::CanMove() const {
162   // It is safe to move the Tensor if and only if we own the unique reference to
163   // it. In that case, we might as well not delete and reallocate, but a future
164   // implementation might need to do so.
165   TensorBuffer* buf = tensorflow::TensorCApi::Buffer(tensor_);
166   if (buf->RefCountIsOne() && buf->root_buffer()->RefCountIsOne() &&
167       buf->OwnsMemory()) {
168     return true;
169   }
170   return false;
171 }
172 
Type() const173 TF_DataType TensorInterface::Type() const {
174   return static_cast<TF_DataType>(tensor_.dtype());
175 }
176 
NumDims() const177 int TensorInterface::NumDims() const { return tensor_.dims(); }
178 
Dim(int dim_index) const179 int64_t TensorInterface::Dim(int dim_index) const {
180   return static_cast<int64_t>(tensor_.dim_size(dim_index));
181 }
182 
NumElements() const183 int64_t TensorInterface::NumElements() const {
184   return static_cast<int64_t>(tensor_.NumElements());
185 }
186 
ByteSize() const187 size_t TensorInterface::ByteSize() const {
188   return tensorflow::TensorCApi::Buffer(tensor_)->size();
189 }
190 
Data() const191 void* TensorInterface::Data() const {
192   return tensorflow::TensorCApi::Buffer(tensor_)->data();
193 }
194 
BitcastFrom(const TensorInterface & from,TF_DataType type,const int64_t * new_dims,int num_new_dims)195 Status TensorInterface::BitcastFrom(const TensorInterface& from,
196                                     TF_DataType type, const int64_t* new_dims,
197                                     int num_new_dims) {
198   tensorflow::TensorShape s;
199   for (int i = 0; i < num_new_dims; ++i) {
200     s.AddDim(new_dims[i]);
201   }
202   return tensor_.BitcastFrom(from.tensor_,
203                              static_cast<tensorflow::DataType>(type), s);
204 }
205 
206 }  // namespace tensorflow
207 
208 // --------------------------------------------------------------------------
StringEncode(const char * src,size_t src_len,char * dst)209 void StringEncode(const char* src, size_t src_len, char* dst) {
210   dst = tensorflow::core::EncodeVarint64(dst, src_len);
211   memcpy(dst, src, src_len);
212 }
213 
TF_StringEncode(const char * src,size_t src_len,char * dst,size_t dst_len,TF_Status * status)214 size_t TF_StringEncode(const char* src, size_t src_len, char* dst,
215                        size_t dst_len, TF_Status* status) {
216   const size_t sz = TF_StringEncodedSize(src_len);
217   if (sz < src_len) {
218     Set_TF_Status_from_Status(
219         status, InvalidArgument("src string is too large to encode"));
220     return 0;
221   }
222   if (dst_len < sz) {
223     Set_TF_Status_from_Status(
224         status,
225         InvalidArgument("dst_len (", dst_len, ") too small to encode a ",
226                         src_len, "-byte string"));
227     return 0;
228   }
229   StringEncode(src, src_len, dst);
230   return sz;
231 }
232 
TF_StringDecode_Impl(const char * src,size_t src_len,const char ** dst,size_t * dst_len)233 static Status TF_StringDecode_Impl(const char* src, size_t src_len,
234                                    const char** dst, size_t* dst_len) {
235   tensorflow::uint64 len64 = 0;
236   const char* p = tensorflow::core::GetVarint64Ptr(src, src + src_len, &len64);
237   if (p == nullptr) {
238     return InvalidArgument("invalid string encoding or truncated src buffer");
239   }
240   if (len64 > std::numeric_limits<size_t>::max()) {
241     return InvalidArgument("encoded string is ", len64,
242                            "-bytes, which is too large for this architecture");
243   }
244   *dst = p;
245   *dst_len = static_cast<size_t>(len64);
246   return Status::OK();
247 }
248 
TF_StringDecode(const char * src,size_t src_len,const char ** dst,size_t * dst_len,TF_Status * status)249 size_t TF_StringDecode(const char* src, size_t src_len, const char** dst,
250                        size_t* dst_len, TF_Status* status) {
251   Set_TF_Status_from_Status(status,
252                             TF_StringDecode_Impl(src, src_len, dst, dst_len));
253   if (TF_GetCode(status) != TF_OK) return 0;
254   return static_cast<size_t>(*dst - src) + *dst_len;
255 }
256 
TF_StringEncodedSize(size_t len)257 size_t TF_StringEncodedSize(size_t len) {
258   return static_cast<size_t>(tensorflow::core::VarintLength(len)) + len;
259 }
260 
DeleteArray(void * data,size_t size,void * arg)261 static void DeleteArray(void* data, size_t size, void* arg) {
262   DCHECK_EQ(data, arg);
263   delete[] reinterpret_cast<char*>(arg);
264 }
265 
266 // Create an empty tensor of type 'dtype'. 'shape' can be arbitrary, but has to
267 // result in a zero-sized tensor.
EmptyTensor(TF_DataType dtype,const tensorflow::TensorShape & shape)268 static TF_Tensor* EmptyTensor(TF_DataType dtype,
269                               const tensorflow::TensorShape& shape) {
270   static char empty;
271   tensorflow::int64 nelems = 1;
272   std::vector<tensorflow::int64> dims;
273   for (int i = 0; i < shape.dims(); ++i) {
274     dims.push_back(shape.dim_size(i));
275     nelems *= shape.dim_size(i);
276   }
277   CHECK_EQ(nelems, 0);
278   static_assert(sizeof(int64_t) == sizeof(tensorflow::int64),
279                 "64-bit int types should match in size");
280   return TF_NewTensor(
281       dtype, reinterpret_cast<const int64_t*>(dims.data()), shape.dims(),
282       reinterpret_cast<void*>(&empty), 0, [](void*, size_t, void*) {}, nullptr);
283 }
284 
285 namespace tensorflow {
286 
287 // Non-static for testing.
TF_TensorFromTensor(const tensorflow::Tensor & src,Status * status)288 TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src, Status* status) {
289   *status = tensorflow::Status::OK();
290   if (!src.IsInitialized()) {
291     *status = FailedPrecondition(
292         "attempt to use a tensor with an uninitialized value");
293     return nullptr;
294   }
295   if (src.NumElements() == 0) {
296     return EmptyTensor(static_cast<TF_DataType>(src.dtype()), src.shape());
297   }
298   if (src.dtype() == tensorflow::DT_RESOURCE) {
299     if (src.shape().dims() != 0) {
300       *status = InvalidArgument(
301           "Unexpected non-scalar DT_RESOURCE tensor seen (shape: ",
302           src.shape().DebugString(),
303           "). Please file a bug at "
304           "https://github.com/tensorflow/tensorflow/issues/new, "
305           "ideally with a "
306           "short code snippet that reproduces this error.");
307       return nullptr;
308     }
309     const string str =
310         src.scalar<tensorflow::ResourceHandle>()().SerializeAsString();
311     TF_Tensor* t = TF_AllocateTensor(TF_RESOURCE, {}, 0, str.size());
312     std::memcpy(TF_TensorData(t), str.c_str(), str.size());
313     return t;
314   }
315   if (src.dtype() != tensorflow::DT_STRING) {
316     Tensor tensor;
317     if (!tensor.CopyFrom(src, src.shape())) {
318       return nullptr;
319     }
320     return new TF_Tensor{std::make_unique<tensorflow::TensorInterface>(tensor)};
321   }
322   // DT_STRING tensors require a copying since TF_Tensor.buffer expects a flatly
323   // encoded sequence of strings.
324 
325   // Compute bytes needed for encoding.
326   size_t size = 0;
327   const auto& srcarray = src.flat<tstring>();
328   for (int i = 0; i < srcarray.size(); ++i) {
329     const string& s = srcarray(i);
330     // uint64 starting_offset, TF_StringEncode-d string.
331     size += sizeof(tensorflow::uint64) + TF_StringEncodedSize(s.size());
332   }
333 
334   // Encode all strings.
335   char* base = new char[size];
336   char* data_start = base + sizeof(tensorflow::uint64) * srcarray.size();
337   char* dst = data_start;  // Where next string is encoded.
338   size_t dst_len = size - static_cast<size_t>(data_start - base);
339   tensorflow::uint64* offsets = reinterpret_cast<tensorflow::uint64*>(base);
340   for (int i = 0; i < srcarray.size(); ++i) {
341     *offsets = (dst - data_start);
342     offsets++;
343     const string& s = srcarray(i);
344     const size_t consumed = TF_StringEncodedSize(s.size());
345     StringEncode(s.data(), s.size(), dst);
346     dst += consumed;
347     dst_len -= consumed;
348   }
349   if (dst != base + size) {
350     *status = InvalidArgument(
351         "invalid string tensor encoding (decoded ", (dst - base),
352         " bytes, but the tensor is encoded in ", size, " bytes");
353     delete[] base;
354     return nullptr;
355   }
356 
357   auto dims = src.shape().dim_sizes();
358   std::vector<tensorflow::int64> dimvec(dims.size());
359   for (size_t i = 0; i < dims.size(); ++i) {
360     dimvec[i] = dims[i];
361   }
362   static_assert(sizeof(int64_t) == sizeof(tensorflow::int64),
363                 "64-bit int types should match in size");
364   return TF_NewTensor(TF_STRING,
365                       reinterpret_cast<const int64_t*>(dimvec.data()),
366                       dimvec.size(), base, size, DeleteArray, base);
367 }
368 
TF_TensorToTensor(const TF_Tensor * src,Tensor * dst)369 Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst) {
370   return static_cast<const tensorflow::TensorInterface*>(src->tensor.get())
371       ->ToTensor(dst);
372 }
373 
ToTensor(Tensor * dst) const374 Status TensorInterface::ToTensor(Tensor* dst) const {
375   if (tensor_.dtype() == DT_RESOURCE) {
376     if (tensor_.dims() != 0) {
377       return InvalidArgument(
378           "Malformed TF_RESOURCE tensor: expected a scalar, got a tensor with "
379           "shape ",
380           tensor_.shape().DebugString());
381     }
382     *dst = Tensor(tensorflow::DT_RESOURCE, tensor_.shape());
383     if (!dst->scalar<tensorflow::ResourceHandle>()().ParseFromString(
384             string(static_cast<const char*>(Data()), ByteSize()))) {
385       return InvalidArgument(
386           "Malformed TF_RESOURCE tensor: unable to parse resource handle");
387     }
388     return Status::OK();
389   }
390   if (tensor_.dtype() != DT_STRING) {
391     *dst = tensor_;
392     return Status::OK();
393   }
394   // TF_STRING tensors require copying since Tensor class expects a sequence of
395   // string objects.
396   const tensorflow::int64 num_elements = tensor_.NumElements();
397   const char* input = reinterpret_cast<const char*>(Data());
398   const size_t src_size = ByteSize();
399   if (static_cast<tensorflow::int64>(src_size / sizeof(tensorflow::uint64)) <
400       num_elements) {
401     return InvalidArgument(
402         "Malformed TF_STRING tensor; too short to hold number of elements");
403   }
404   const char* data_start = input + sizeof(tensorflow::uint64) * num_elements;
405   const char* limit = input + src_size;
406 
407   *dst = Tensor(tensor_.dtype(), tensor_.shape());
408   auto dstarray = dst->flat<tstring>();
409   for (tensorflow::int64 i = 0; i < num_elements; ++i) {
410     tensorflow::uint64 offset =
411         reinterpret_cast<const tensorflow::uint64*>(input)[i];
412     if (static_cast<ptrdiff_t>(offset) >= (limit - data_start)) {
413       return InvalidArgument("Malformed TF_STRING tensor; element ", i,
414                              " out of range");
415     }
416     size_t len;
417     const char* p;
418     const char* srcp = data_start + offset;
419     Status status = TF_StringDecode_Impl(srcp, limit - srcp, &p, &len);
420     if (!status.ok()) return status;
421     dstarray(i).assign(p, len);
422   }
423   return Status::OK();
424 }
425 
IsAligned() const426 bool TensorInterface::IsAligned() const { return tensor_.IsAligned(); }
427 
428 }  // namespace tensorflow
429 
TF_TensorIsAligned(const TF_Tensor * t)430 bool TF_TensorIsAligned(const TF_Tensor* t) { return t->tensor->IsAligned(); }
431