1 /* Copyright 2016 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/python/lib/core/ndarray_tensor.h"
17
18 #include <cstring>
19
20 #include "tensorflow/core/lib/core/coding.h"
21 #include "tensorflow/core/lib/core/errors.h"
22 #include "tensorflow/core/lib/gtl/inlined_vector.h"
23 #include "tensorflow/core/platform/types.h"
24 #include "tensorflow/python/lib/core/bfloat16.h"
25 #include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
26
27 namespace tensorflow {
28 namespace {
29
PyArrayDescr_to_TF_DataType(PyArray_Descr * descr,TF_DataType * out_tf_datatype)30 Status PyArrayDescr_to_TF_DataType(PyArray_Descr* descr,
31 TF_DataType* out_tf_datatype) {
32 PyObject* key;
33 PyObject* value;
34 Py_ssize_t pos = 0;
35 if (PyDict_Next(descr->fields, &pos, &key, &value)) {
36 // In Python 3, the keys of numpy custom struct types are unicode, unlike
37 // Python 2, where the keys are bytes.
38 const char* key_string =
39 PyBytes_Check(key) ? PyBytes_AsString(key)
40 : PyBytes_AsString(PyUnicode_AsASCIIString(key));
41 if (!key_string) {
42 return errors::Internal("Corrupt numpy type descriptor");
43 }
44 tensorflow::string key = key_string;
45 // The typenames here should match the field names in the custom struct
46 // types constructed in test_util.py.
47 // TODO(mrry,keveman): Investigate Numpy type registration to replace this
48 // hard-coding of names.
49 if (key == "quint8") {
50 *out_tf_datatype = TF_QUINT8;
51 } else if (key == "qint8") {
52 *out_tf_datatype = TF_QINT8;
53 } else if (key == "qint16") {
54 *out_tf_datatype = TF_QINT16;
55 } else if (key == "quint16") {
56 *out_tf_datatype = TF_QUINT16;
57 } else if (key == "qint32") {
58 *out_tf_datatype = TF_QINT32;
59 } else if (key == "resource") {
60 *out_tf_datatype = TF_RESOURCE;
61 } else {
62 return errors::Internal("Unsupported numpy data type");
63 }
64 return Status::OK();
65 }
66 return errors::Internal("Unsupported numpy data type");
67 }
68
PyArray_TYPE_to_TF_DataType(PyArrayObject * array,TF_DataType * out_tf_datatype)69 Status PyArray_TYPE_to_TF_DataType(PyArrayObject* array,
70 TF_DataType* out_tf_datatype) {
71 int pyarray_type = PyArray_TYPE(array);
72 PyArray_Descr* descr = PyArray_DESCR(array);
73 switch (pyarray_type) {
74 case NPY_FLOAT16:
75 *out_tf_datatype = TF_HALF;
76 break;
77 case NPY_FLOAT32:
78 *out_tf_datatype = TF_FLOAT;
79 break;
80 case NPY_FLOAT64:
81 *out_tf_datatype = TF_DOUBLE;
82 break;
83 case NPY_INT32:
84 *out_tf_datatype = TF_INT32;
85 break;
86 case NPY_UINT8:
87 *out_tf_datatype = TF_UINT8;
88 break;
89 case NPY_UINT16:
90 *out_tf_datatype = TF_UINT16;
91 break;
92 case NPY_UINT32:
93 *out_tf_datatype = TF_UINT32;
94 break;
95 case NPY_UINT64:
96 *out_tf_datatype = TF_UINT64;
97 break;
98 case NPY_INT8:
99 *out_tf_datatype = TF_INT8;
100 break;
101 case NPY_INT16:
102 *out_tf_datatype = TF_INT16;
103 break;
104 case NPY_INT64:
105 *out_tf_datatype = TF_INT64;
106 break;
107 case NPY_BOOL:
108 *out_tf_datatype = TF_BOOL;
109 break;
110 case NPY_COMPLEX64:
111 *out_tf_datatype = TF_COMPLEX64;
112 break;
113 case NPY_COMPLEX128:
114 *out_tf_datatype = TF_COMPLEX128;
115 break;
116 case NPY_OBJECT:
117 case NPY_STRING:
118 case NPY_UNICODE:
119 *out_tf_datatype = TF_STRING;
120 break;
121 case NPY_VOID:
122 // Quantized types are currently represented as custom struct types.
123 // PyArray_TYPE returns NPY_VOID for structs, and we should look into
124 // descr to derive the actual type.
125 // Direct feeds of certain types of ResourceHandles are represented as a
126 // custom struct type.
127 return PyArrayDescr_to_TF_DataType(descr, out_tf_datatype);
128 default:
129 if (pyarray_type == Bfloat16NumpyType()) {
130 *out_tf_datatype = TF_BFLOAT16;
131 break;
132 }
133 // TODO(mrry): Support these.
134 return errors::Internal("Unsupported feed type");
135 }
136 return Status::OK();
137 }
138
PyObjectToString(PyObject * obj,const char ** ptr,Py_ssize_t * len,PyObject ** ptr_owner)139 Status PyObjectToString(PyObject* obj, const char** ptr, Py_ssize_t* len,
140 PyObject** ptr_owner) {
141 *ptr_owner = nullptr;
142 if (!PyUnicode_Check(obj)) {
143 char* buf;
144 if (PyBytes_AsStringAndSize(obj, &buf, len) != 0) {
145 return errors::Internal("Unable to get element as bytes.");
146 }
147 *ptr = buf;
148 return Status::OK();
149 }
150 #if (PY_MAJOR_VERSION > 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION >= 3))
151 *ptr = PyUnicode_AsUTF8AndSize(obj, len);
152 if (*ptr != nullptr) return Status::OK();
153 #else
154 PyObject* utemp = PyUnicode_AsUTF8String(obj);
155 char* buf;
156 if (utemp != nullptr && PyBytes_AsStringAndSize(utemp, &buf, len) != -1) {
157 *ptr = buf;
158 *ptr_owner = utemp;
159 return Status::OK();
160 }
161 Py_XDECREF(utemp);
162 #endif
163 return errors::Internal("Unable to convert element to UTF-8.");
164 }
165
166 // Iterate over the string array 'array', extract the ptr and len of each string
167 // element and call f(ptr, len).
168 template <typename F>
PyBytesArrayMap(PyArrayObject * array,F f)169 Status PyBytesArrayMap(PyArrayObject* array, F f) {
170 Safe_PyObjectPtr iter = tensorflow::make_safe(
171 PyArray_IterNew(reinterpret_cast<PyObject*>(array)));
172 while (PyArray_ITER_NOTDONE(iter.get())) {
173 auto item = tensorflow::make_safe(PyArray_GETITEM(
174 array, static_cast<char*>(PyArray_ITER_DATA(iter.get()))));
175 if (!item) {
176 return errors::Internal("Unable to get element from the feed - no item.");
177 }
178 Py_ssize_t len;
179 const char* ptr;
180 PyObject* ptr_owner;
181 TF_RETURN_IF_ERROR(PyObjectToString(item.get(), &ptr, &len, &ptr_owner));
182 f(ptr, len);
183 Py_XDECREF(ptr_owner);
184 PyArray_ITER_NEXT(iter.get());
185 }
186 return Status::OK();
187 }
188
189 // Encode the strings in 'array' into a contiguous buffer and return the base of
190 // the buffer. The caller takes ownership of the buffer.
EncodePyBytesArray(PyArrayObject * array,tensorflow::int64 nelems,size_t * size,void ** buffer)191 Status EncodePyBytesArray(PyArrayObject* array, tensorflow::int64 nelems,
192 size_t* size, void** buffer) {
193 // Compute bytes needed for encoding.
194 *size = 0;
195 TF_RETURN_IF_ERROR(
196 PyBytesArrayMap(array, [&size](const char* ptr, Py_ssize_t len) {
197 *size += sizeof(tensorflow::uint64) +
198 tensorflow::core::VarintLength(len) + len;
199 }));
200 // Encode all strings.
201 std::unique_ptr<char[]> base_ptr(new char[*size]);
202 char* base = base_ptr.get();
203 char* data_start = base + sizeof(tensorflow::uint64) * nelems;
204 char* dst = data_start; // Where next string is encoded.
205 tensorflow::uint64* offsets = reinterpret_cast<tensorflow::uint64*>(base);
206
207 TF_RETURN_IF_ERROR(PyBytesArrayMap(
208 array, [&data_start, &dst, &offsets](const char* ptr, Py_ssize_t len) {
209 *offsets = (dst - data_start);
210 offsets++;
211 dst = tensorflow::core::EncodeVarint64(dst, len);
212 memcpy(dst, ptr, len);
213 dst += len;
214 }));
215 CHECK_EQ(dst, base + *size);
216 *buffer = base_ptr.release();
217 return Status::OK();
218 }
219
CopyTF_TensorStringsToPyArray(const TF_Tensor * src,uint64 nelems,PyArrayObject * dst)220 Status CopyTF_TensorStringsToPyArray(const TF_Tensor* src, uint64 nelems,
221 PyArrayObject* dst) {
222 const void* tensor_data = TF_TensorData(src);
223 const size_t tensor_size = TF_TensorByteSize(src);
224 const char* limit = static_cast<const char*>(tensor_data) + tensor_size;
225 DCHECK(tensor_data != nullptr);
226 DCHECK_EQ(TF_STRING, TF_TensorType(src));
227
228 const uint64* offsets = static_cast<const uint64*>(tensor_data);
229 const size_t offsets_size = sizeof(uint64) * nelems;
230 const char* data = static_cast<const char*>(tensor_data) + offsets_size;
231
232 const size_t expected_tensor_size =
233 (limit - static_cast<const char*>(tensor_data));
234 if (expected_tensor_size - tensor_size) {
235 return errors::InvalidArgument(
236 "Invalid/corrupt TF_STRING tensor: expected ", expected_tensor_size,
237 " bytes of encoded strings for the tensor containing ", nelems,
238 " strings, but the tensor is encoded in ", tensor_size, " bytes");
239 }
240 std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
241 TF_NewStatus(), TF_DeleteStatus);
242 auto iter = make_safe(PyArray_IterNew(reinterpret_cast<PyObject*>(dst)));
243 for (int64 i = 0; i < nelems; ++i) {
244 const char* start = data + offsets[i];
245 const char* ptr = nullptr;
246 size_t len = 0;
247
248 TF_StringDecode(start, limit - start, &ptr, &len, status.get());
249 if (TF_GetCode(status.get()) != TF_OK) {
250 return errors::InvalidArgument(TF_Message(status.get()));
251 }
252
253 auto py_string = make_safe(PyBytes_FromStringAndSize(ptr, len));
254 if (py_string == nullptr) {
255 return errors::Internal(
256 "failed to create a python byte array when converting element #", i,
257 " of a TF_STRING tensor to a numpy ndarray");
258 }
259
260 if (PyArray_SETITEM(dst, static_cast<char*>(PyArray_ITER_DATA(iter.get())),
261 py_string.get()) != 0) {
262 return errors::Internal("Error settings element #", i,
263 " in the numpy ndarray");
264 }
265 PyArray_ITER_NEXT(iter.get());
266 }
267 return Status::OK();
268 }
269
270 // Determine the dimensions of a numpy ndarray to be created to represent an
271 // output Tensor.
GetPyArrayDimensionsForTensor(const TF_Tensor * tensor,tensorflow::int64 * nelems)272 gtl::InlinedVector<npy_intp, 4> GetPyArrayDimensionsForTensor(
273 const TF_Tensor* tensor, tensorflow::int64* nelems) {
274 const int ndims = TF_NumDims(tensor);
275 gtl::InlinedVector<npy_intp, 4> dims(ndims);
276 if (TF_TensorType(tensor) == TF_RESOURCE) {
277 CHECK_EQ(ndims, 0)
278 << "Fetching of non-scalar resource tensors is not supported.";
279 dims.push_back(TF_TensorByteSize(tensor));
280 *nelems = dims[0];
281 } else {
282 *nelems = 1;
283 for (int i = 0; i < ndims; ++i) {
284 dims[i] = TF_Dim(tensor, i);
285 *nelems *= dims[i];
286 }
287 }
288 return dims;
289 }
290
291 // Determine the type description (PyArray_Descr) of a numpy ndarray to be
292 // created to represent an output Tensor.
GetPyArrayDescrForTensor(const TF_Tensor * tensor,PyArray_Descr ** descr)293 Status GetPyArrayDescrForTensor(const TF_Tensor* tensor,
294 PyArray_Descr** descr) {
295 if (TF_TensorType(tensor) == TF_RESOURCE) {
296 PyObject* field = PyTuple_New(3);
297 #if PY_MAJOR_VERSION < 3
298 PyTuple_SetItem(field, 0, PyBytes_FromString("resource"));
299 #else
300 PyTuple_SetItem(field, 0, PyUnicode_FromString("resource"));
301 #endif
302 PyTuple_SetItem(field, 1, PyArray_TypeObjectFromType(NPY_UBYTE));
303 PyTuple_SetItem(field, 2, PyLong_FromLong(1));
304 PyObject* fields = PyList_New(1);
305 PyList_SetItem(fields, 0, field);
306 int convert_result = PyArray_DescrConverter(fields, descr);
307 Py_CLEAR(field);
308 Py_CLEAR(fields);
309 if (convert_result != 1) {
310 return errors::Internal("Failed to create numpy array description for ",
311 "TF_RESOURCE-type tensor");
312 }
313 } else {
314 int type_num = -1;
315 TF_RETURN_IF_ERROR(
316 TF_DataType_to_PyArray_TYPE(TF_TensorType(tensor), &type_num));
317 *descr = PyArray_DescrFromType(type_num);
318 }
319
320 return Status::OK();
321 }
322
FastMemcpy(void * dst,const void * src,size_t size)323 inline void FastMemcpy(void* dst, const void* src, size_t size) {
324 // clang-format off
325 switch (size) {
326 // Most compilers will generate inline code for fixed sizes,
327 // which is significantly faster for small copies.
328 case 1: memcpy(dst, src, 1); break;
329 case 2: memcpy(dst, src, 2); break;
330 case 3: memcpy(dst, src, 3); break;
331 case 4: memcpy(dst, src, 4); break;
332 case 5: memcpy(dst, src, 5); break;
333 case 6: memcpy(dst, src, 6); break;
334 case 7: memcpy(dst, src, 7); break;
335 case 8: memcpy(dst, src, 8); break;
336 case 9: memcpy(dst, src, 9); break;
337 case 10: memcpy(dst, src, 10); break;
338 case 11: memcpy(dst, src, 11); break;
339 case 12: memcpy(dst, src, 12); break;
340 case 13: memcpy(dst, src, 13); break;
341 case 14: memcpy(dst, src, 14); break;
342 case 15: memcpy(dst, src, 15); break;
343 case 16: memcpy(dst, src, 16); break;
344 #if defined(PLATFORM_GOOGLE) || defined(PLATFORM_POSIX) && \
345 !defined(IS_MOBILE_PLATFORM)
346 // On Linux, memmove appears to be faster than memcpy for
347 // large sizes, strangely enough.
348 default: memmove(dst, src, size); break;
349 #else
350 default: memcpy(dst, src, size); break;
351 #endif
352 }
353 // clang-format on
354 }
355
356 } // namespace
357
358 // Converts the given TF_Tensor to a numpy ndarray.
359 // If the returned status is OK, the caller becomes the owner of *out_array.
TF_TensorToPyArray(Safe_TF_TensorPtr tensor,PyObject ** out_ndarray)360 Status TF_TensorToPyArray(Safe_TF_TensorPtr tensor, PyObject** out_ndarray) {
361 // A fetched operation will correspond to a null tensor, and a None
362 // in Python.
363 if (tensor == nullptr) {
364 Py_INCREF(Py_None);
365 *out_ndarray = Py_None;
366 return Status::OK();
367 }
368 int64 nelems = -1;
369 gtl::InlinedVector<npy_intp, 4> dims =
370 GetPyArrayDimensionsForTensor(tensor.get(), &nelems);
371
372 // If the type is neither string nor resource we can reuse the Tensor memory.
373 TF_Tensor* original = tensor.get();
374 TF_Tensor* moved = TF_TensorMaybeMove(tensor.release());
375 if (moved != nullptr) {
376 if (ArrayFromMemory(dims.size(), dims.data(), TF_TensorData(moved),
377 static_cast<DataType>(TF_TensorType(moved)),
378 [moved] { TF_DeleteTensor(moved); }, out_ndarray)
379 .ok()) {
380 return Status::OK();
381 }
382 }
383 tensor.reset(original);
384
385 // Copy the TF_TensorData into a newly-created ndarray and return it.
386 PyArray_Descr* descr = nullptr;
387 TF_RETURN_IF_ERROR(GetPyArrayDescrForTensor(tensor.get(), &descr));
388 Safe_PyObjectPtr safe_out_array =
389 tensorflow::make_safe(PyArray_Empty(dims.size(), dims.data(), descr, 0));
390 if (!safe_out_array) {
391 return errors::Internal("Could not allocate ndarray");
392 }
393 PyArrayObject* py_array =
394 reinterpret_cast<PyArrayObject*>(safe_out_array.get());
395 if (TF_TensorType(tensor.get()) == TF_STRING) {
396 Status s = CopyTF_TensorStringsToPyArray(tensor.get(), nelems, py_array);
397 if (!s.ok()) {
398 return s;
399 }
400 } else if (static_cast<size_t>(PyArray_NBYTES(py_array)) !=
401 TF_TensorByteSize(tensor.get())) {
402 return errors::Internal("ndarray was ", PyArray_NBYTES(py_array),
403 " bytes but TF_Tensor was ",
404 TF_TensorByteSize(tensor.get()), " bytes");
405 } else {
406 FastMemcpy(PyArray_DATA(py_array), TF_TensorData(tensor.get()),
407 PyArray_NBYTES(py_array));
408 }
409
410 // PyArray_Return turns rank 0 arrays into numpy scalars
411 *out_ndarray = PyArray_Return(
412 reinterpret_cast<PyArrayObject*>(safe_out_array.release()));
413 return Status::OK();
414 }
415
PyArrayToTF_Tensor(PyObject * ndarray,Safe_TF_TensorPtr * out_tensor)416 Status PyArrayToTF_Tensor(PyObject* ndarray, Safe_TF_TensorPtr* out_tensor) {
417 DCHECK(out_tensor != nullptr);
418
419 // Make sure we dereference this array object in case of error, etc.
420 Safe_PyObjectPtr array_safe(make_safe(
421 PyArray_FromAny(ndarray, nullptr, 0, 0, NPY_ARRAY_CARRAY_RO, nullptr)));
422 if (!array_safe) return errors::InvalidArgument("Not a ndarray.");
423 PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());
424
425 // Convert numpy dtype to TensorFlow dtype.
426 TF_DataType dtype = TF_FLOAT;
427 TF_RETURN_IF_ERROR(PyArray_TYPE_to_TF_DataType(array, &dtype));
428
429 tensorflow::int64 nelems = 1;
430 gtl::InlinedVector<int64_t, 4> dims;
431 for (int i = 0; i < PyArray_NDIM(array); ++i) {
432 dims.push_back(PyArray_SHAPE(array)[i]);
433 nelems *= dims[i];
434 }
435
436 // Create a TF_Tensor based on the fed data. In the case of non-string data
437 // type, this steals a reference to array, which will be relinquished when
438 // the underlying buffer is deallocated. For string, a new temporary buffer
439 // is allocated into which the strings are encoded.
440 if (dtype == TF_RESOURCE) {
441 size_t size = PyArray_NBYTES(array);
442 array_safe.release();
443 *out_tensor = make_safe(TF_NewTensor(dtype, {}, 0, PyArray_DATA(array),
444 size, &DelayedNumpyDecref, array));
445
446 } else if (dtype != TF_STRING) {
447 size_t size = PyArray_NBYTES(array);
448 array_safe.release();
449 *out_tensor = make_safe(TF_NewTensor(dtype, dims.data(), dims.size(),
450 PyArray_DATA(array), size,
451 &DelayedNumpyDecref, array));
452 } else {
453 size_t size = 0;
454 void* encoded = nullptr;
455 TF_RETURN_IF_ERROR(EncodePyBytesArray(array, nelems, &size, &encoded));
456 *out_tensor =
457 make_safe(TF_NewTensor(dtype, dims.data(), dims.size(), encoded, size,
458 [](void* data, size_t len, void* arg) {
459 delete[] reinterpret_cast<char*>(data);
460 },
461 nullptr));
462 }
463 return Status::OK();
464 }
465
466 Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
467 TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src,
468 TF_Status* status);
469
NdarrayToTensor(PyObject * obj,Tensor * ret)470 Status NdarrayToTensor(PyObject* obj, Tensor* ret) {
471 Safe_TF_TensorPtr tf_tensor = make_safe(static_cast<TF_Tensor*>(nullptr));
472 Status s = PyArrayToTF_Tensor(obj, &tf_tensor);
473 if (!s.ok()) {
474 return s;
475 }
476 return TF_TensorToTensor(tf_tensor.get(), ret);
477 }
478
TensorToNdarray(const Tensor & t,PyObject ** ret)479 Status TensorToNdarray(const Tensor& t, PyObject** ret) {
480 TF_Status* status = TF_NewStatus();
481 Safe_TF_TensorPtr tf_tensor = make_safe(TF_TensorFromTensor(t, status));
482 Status tf_status = StatusFromTF_Status(status);
483 TF_DeleteStatus(status);
484 if (!tf_status.ok()) {
485 return tf_status;
486 }
487 return TF_TensorToPyArray(std::move(tf_tensor), ret);
488 }
489
490 } // namespace tensorflow
491