From 0b3a25d684ed2466d35661683baadc7eecce73d3 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Mon, 31 Jul 2017 14:51:11 -0700 Subject: [PATCH] Make TF_RESOURCE memory handling consistent with other types. TF_Tensor's are backed by a contiguous memory region for all but TF_RESOURCE tensors. The memory management of TF_RESOURCE tensors required keeping a backing tensorflow::ResourceHandle* object alive for the lifetime of the TF_Tensor object. This change removes that discrepancy, making the memory backing TF_RESOURCE tensors self-contained. This simplifies use of TF_RESOURCE tensors in the C API (as users of the C API do not need to worry about a tensorflow::ResourceHandle object and its lifetime). In doing so, this moves a string memory copy from the TF_Tensor <-> Numpy conversion to the C API from the Python session helper. Unfortunately, I couldn't figure out how to add a simple unittest in c_api_test.cc. The more comprehensive tensorflow/python/kernel_tests/session_ops_test.py does cover the changed lines though. Additionally, avoid an unnecessary copy when creating TF_STRING or TF_RESOURCE tensors (as eigen alignment is not a requirement for them). PiperOrigin-RevId: 163751880 --- tensorflow/c/c_api.cc | 52 +++++++++- .../core/common_runtime/direct_session.cc | 9 +- tensorflow/python/client/tf_session_helper.cc | 94 ++++++------------- 3 files changed, 84 insertions(+), 71 deletions(-) diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index f5013c3f6a1..570e84ecc24 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -37,6 +37,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -219,9 +220,16 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, TF_ManagedBuffer* buf = new TF_ManagedBuffer; buf->len_ = len; - if (reinterpret_cast(data) % EIGEN_MAX_ALIGN_BYTES != 0) { - // Copy the data into a buffer that satisfies Eigen's alignment - // requirements. + if (dtype != TF_STRING && dtype != TF_RESOURCE && + tensorflow::DataTypeCanUseMemcpy(static_cast(dtype)) && + reinterpret_cast(data) % EIGEN_MAX_ALIGN_BYTES != 0) { + // TF_STRING and TF_RESOURCE tensors have a different representation in + // TF_Tensor than they do in tensorflow::Tensor. So a copy here is a waste + // (any alignement requirements will be taken care of by TF_TensorToTensor + // and TF_TensorFromTensor). + // + // Other types have the same represntation, so copy only if it is safe to do + // so. buf->data_ = allocate_tensor("TF_NewTensor", len); std::memcpy(buf->data_, data, len); buf->deallocator_ = deallocate_buffer; @@ -419,7 +427,30 @@ void TF_Reset(const TF_SessionOptions* opt, const char** containers, namespace tensorflow { Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst) { + if (src->dtype == TF_RESOURCE) { + if (src->buffer->device() != nullptr) { + return InvalidArgument( + "TF_RESOURCE tensor must be placed in host memory"); + } + if (src->shape.dims() != 0) { + return InvalidArgument( + "Malformed TF_RESOURCE tensor: expected a scalar, got a tensor with " + "shape ", + src->shape.DebugString()); + } + *dst = Tensor(DT_RESOURCE, src->shape); + if (!dst->scalar()().ParseFromString( + string(static_cast(TF_TensorData(src)), + TF_TensorByteSize(src)))) { + return InvalidArgument( + "Malformed TF_RESOUCE tensor: unable to parse resource handle"); + } + return Status::OK(); + } if (src->dtype != TF_STRING) { + if (src->buffer->device() != nullptr) { + return InvalidArgument("TF_STRING tensor must be placed in host memory"); + } *dst = TensorCApi::MakeTensor(src->dtype, src->shape, src->buffer->buffer()); return Status::OK(); @@ -458,6 +489,21 @@ Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst) { // Non-static for testing. TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src) { + if (src.dtype() == DT_RESOURCE) { + DCHECK_EQ(0, src.shape().dims()) << src.shape().DebugString(); + if (src.shape().dims() != 0) { + LOG(ERROR) << "Unexpected non-scalar DT_RESOURCE tensor seen (shape: " + << src.shape().DebugString() + << "). Please file a bug at " + "https://github.com/tensorflow/tensorflow/issues/new, " + "ideally with a " + "short code snippet that reproduces this error."; + } + const string str = src.scalar()().SerializeAsString(); + TF_Tensor* t = TF_AllocateTensor(TF_RESOURCE, {}, 0, str.size()); + std::memcpy(TF_TensorData(t), str.c_str(), str.size()); + return t; + } if (src.dtype() != DT_STRING) { TensorBuffer* buf = TensorCApi::Buffer(src); buf->Ref(); diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc index 9b8744dccd7..294ad13a0a4 100644 --- a/tensorflow/core/common_runtime/direct_session.cc +++ b/tensorflow/core/common_runtime/direct_session.cc @@ -877,7 +877,8 @@ Status DirectSession::ResourceHandleToInputTensor(const Tensor& resource_tensor, resource_tensor.dtype())); } - ResourceHandle resource_handle = resource_tensor.scalar()(); + const ResourceHandle& resource_handle = + resource_tensor.scalar()(); if (resource_handle.container() == SessionState::kTensorHandleResourceTypeName) { @@ -886,7 +887,11 @@ Status DirectSession::ResourceHandleToInputTensor(const Tensor& resource_tensor, return errors::InvalidArgument(strings::StrCat( "Invalid resource type hash code: ", resource_handle.hash_code(), "(name: ", resource_handle.name(), - " type: ", resource_handle.maybe_type_name(), ")")); + " type: ", resource_handle.maybe_type_name(), + "). Perhaps a resource tensor was being provided as a feed? That is " + "not currently allowed. Please file an issue at " + "https://github.com/tensorflow/tensorflow/issues/new, ideally with a " + "short code snippet that leads to this error message.")); } } diff --git a/tensorflow/python/client/tf_session_helper.cc b/tensorflow/python/client/tf_session_helper.cc index 5c04a8c2a5e..7935924f238 100644 --- a/tensorflow/python/client/tf_session_helper.cc +++ b/tensorflow/python/client/tf_session_helper.cc @@ -255,25 +255,19 @@ static Status CopyStringToPyArrayElement(PyArrayObject* pyarray, void* i_ptr, // output Tensor. gtl::InlinedVector GetPyArrayDimensionsForTensor( const TF_Tensor* tensor, tensorflow::int64* nelems) { + const int ndims = TF_NumDims(tensor); + gtl::InlinedVector dims(ndims); if (TF_TensorType(tensor) == TF_RESOURCE) { - gtl::InlinedVector dims(1); - ResourceHandle* resource_handle = - reinterpret_cast(TF_TensorData(tensor)); - dims[0] = resource_handle->SerializeAsString().size(); + dims[0] = TF_TensorByteSize(tensor); *nelems = dims[0]; - - return dims; } else { - const int ndims = TF_NumDims(tensor); - gtl::InlinedVector dims(ndims); *nelems = 1; for (int i = 0; i < ndims; ++i) { dims[i] = TF_Dim(tensor, i); *nelems *= dims[i]; } - - return dims; } + return dims; } // Determine the type description (PyArray_Descr) of a numpy ndarray to be @@ -348,30 +342,22 @@ Status TFTensorToPyArray(Safe_TF_TensorPtr tensor, PyObject** out_ndarray) { } PyArrayObject* py_array = reinterpret_cast(safe_out_array.get()); - if (PyArray_NBYTES(py_array) != - static_cast(TF_TensorByteSize(tensor.get()))) { - if (TF_TensorType(tensor.get()) == TF_STRING) { - // Copy element by element. - auto iter = tensorflow::make_safe(PyArray_IterNew(safe_out_array.get())); - for (tensorflow::int64 i = 0; i < nelems; ++i) { - auto s = CopyStringToPyArrayElement(py_array, iter.get(), tensor.get(), - nelems, i); - if (!s.ok()) { - return s; - } - PyArray_ITER_NEXT(iter.get()); + if (TF_TensorType(tensor.get()) == TF_STRING) { + // Copy element by element. + auto iter = tensorflow::make_safe(PyArray_IterNew(safe_out_array.get())); + for (tensorflow::int64 i = 0; i < nelems; ++i) { + auto s = CopyStringToPyArrayElement(py_array, iter.get(), tensor.get(), + nelems, i); + if (!s.ok()) { + return s; } - } else if (TF_TensorType(tensor.get()) == TF_RESOURCE) { - ResourceHandle* resource_handle = - reinterpret_cast(TF_TensorData(tensor.get())); - memcpy(PyArray_DATA(py_array), - resource_handle->SerializeAsString().c_str(), - PyArray_NBYTES(py_array)); - } else { - return errors::Internal("ndarray was ", PyArray_NBYTES(py_array), - " bytes but TF_Tensor was ", - TF_TensorByteSize(tensor.get()), " bytes"); + PyArray_ITER_NEXT(iter.get()); } + } else if (static_cast(PyArray_NBYTES(py_array)) != + TF_TensorByteSize(tensor.get())) { + return errors::Internal("ndarray was ", PyArray_NBYTES(py_array), + " bytes but TF_Tensor was ", + TF_TensorByteSize(tensor.get()), " bytes"); } else { memcpy(PyArray_DATA(py_array), TF_TensorData(tensor.get()), PyArray_NBYTES(py_array)); @@ -388,18 +374,9 @@ Status TFTensorToPyArray(Safe_TF_TensorPtr tensor, PyObject** out_ndarray) { // data. After `out_tensor` is destroyed, this reference must (eventually) be // decremented via ClearDecrefCache(). // -// If `ndarray` contains a resource handle, `*resource_handle` will be set to -// the deserialized handle. Otherwise it is set to nullptr. Caller becomes owner -// of `*resource_handle` if it's set, and it must outlive the returned -// `out_tensor`. -// -// `resource_handle` and `out_tensor` must be non-null. Caller retains ownership -// of `ndarray`. -Status PyArrayToTFTensor(PyObject* ndarray, Safe_TF_TensorPtr* out_tensor, - ResourceHandle** resource_handle) { +// `out_tensor` must be non-null. Caller retains ownership of `ndarray`. +Status PyArrayToTFTensor(PyObject* ndarray, Safe_TF_TensorPtr* out_tensor) { DCHECK(out_tensor != nullptr); - DCHECK(resource_handle != nullptr); - *resource_handle = nullptr; // Make sure we dereference this array object in case of error, etc. Safe_PyObjectPtr array_safe(make_safe( @@ -423,16 +400,11 @@ Status PyArrayToTFTensor(PyObject* ndarray, Safe_TF_TensorPtr* out_tensor, // the underlying buffer is deallocated. For string, a new temporary buffer // is allocated into which the strings are encoded. if (dtype == TF_RESOURCE) { - const string serialized(reinterpret_cast(PyArray_DATA(array)), - PyArray_NBYTES(array)); - *resource_handle = new ResourceHandle(); - (*resource_handle)->ParseFromString(serialized); - TF_Tensor* tf_tensor = - TF_AllocateTensor(dtype, {}, 0, sizeof(ResourceHandle)); - std::memcpy(TF_TensorData(tf_tensor), - reinterpret_cast(*resource_handle), - sizeof(ResourceHandle)); - *out_tensor = make_safe(tf_tensor); + size_t size = PyArray_NBYTES(array); + array_safe.release(); + *out_tensor = make_safe(TF_NewTensor(dtype, {}, 0, PyArray_DATA(array), + size, &DelayedNumpyDecref, array)); + } else if (dtype != TF_STRING) { size_t size = PyArray_NBYTES(array); array_safe.release(); @@ -478,7 +450,7 @@ void TF_Run_wrapper_helper(TF_DeprecatedSession* session, const char* handle, NameVector input_names; std::vector inputs_safe; // Used to delete tensors. - TF_TensorVector inputs_unsafe; // Used to contain the arg to TF_Run. + TF_TensorVector inputs_unsafe; // Used to contain the arg to TF_Run. PyObject* key; PyObject* value; @@ -486,7 +458,6 @@ void TF_Run_wrapper_helper(TF_DeprecatedSession* session, const char* handle, int index = 0; Status s; - gtl::InlinedVector, 4> resource_handles; while (PyDict_Next(feed_dict, &pos, &key, &value)) { char* key_string = PyBytes_AsString(key); if (!key_string) { @@ -497,16 +468,12 @@ void TF_Run_wrapper_helper(TF_DeprecatedSession* session, const char* handle, input_names.push_back(key_string); inputs_safe.emplace_back(make_safe(static_cast(nullptr))); - ResourceHandle* resource_handle; - s = PyArrayToTFTensor(value, &inputs_safe.back(), &resource_handle); + s = PyArrayToTFTensor(value, &inputs_safe.back()); if (!s.ok()) { Set_TF_Status_from_Status(out_status, s); return; } inputs_unsafe.push_back(inputs_safe.back().get()); - if (resource_handle != nullptr) { - resource_handles.emplace_back(resource_handle); - } ++index; } @@ -647,14 +614,9 @@ void TF_SessionRun_wrapper_helper(TF_Session* session, const char* handle, // queued or assigned to a variable). TF_TensorVector input_vals; std::vector input_vals_safe; - gtl::InlinedVector, 4> resource_handles; for (PyObject* ndarray : input_ndarrays) { input_vals_safe.emplace_back(make_safe(static_cast(nullptr))); - ResourceHandle* resource_handle; - s = PyArrayToTFTensor(ndarray, &input_vals_safe.back(), &resource_handle); - if (resource_handle != nullptr) { - resource_handles.emplace_back(resource_handle); - } + s = PyArrayToTFTensor(ndarray, &input_vals_safe.back()); if (!s.ok()) { Set_TF_Status_from_Status(out_status, s); return;