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
This commit is contained in:
Asim Shankar 2017-07-31 14:51:11 -07:00 committed by TensorFlower Gardener
parent 1333e77450
commit 0b3a25d684
3 changed files with 84 additions and 71 deletions

View File

@ -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<intptr_t>(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<DataType>(dtype)) &&
reinterpret_cast<intptr_t>(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<ResourceHandle>()().ParseFromString(
string(static_cast<const char*>(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<ResourceHandle>()().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();

View File

@ -877,7 +877,8 @@ Status DirectSession::ResourceHandleToInputTensor(const Tensor& resource_tensor,
resource_tensor.dtype()));
}
ResourceHandle resource_handle = resource_tensor.scalar<ResourceHandle>()();
const ResourceHandle& resource_handle =
resource_tensor.scalar<ResourceHandle>()();
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."));
}
}

View File

@ -255,25 +255,19 @@ static Status CopyStringToPyArrayElement(PyArrayObject* pyarray, void* i_ptr,
// output Tensor.
gtl::InlinedVector<npy_intp, 4> GetPyArrayDimensionsForTensor(
const TF_Tensor* tensor, tensorflow::int64* nelems) {
const int ndims = TF_NumDims(tensor);
gtl::InlinedVector<npy_intp, 4> dims(ndims);
if (TF_TensorType(tensor) == TF_RESOURCE) {
gtl::InlinedVector<npy_intp, 4> dims(1);
ResourceHandle* resource_handle =
reinterpret_cast<ResourceHandle*>(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<npy_intp, 4> 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<PyArrayObject*>(safe_out_array.get());
if (PyArray_NBYTES(py_array) !=
static_cast<int64>(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<ResourceHandle*>(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<size_t>(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<char*>(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<void*>(*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<Safe_TF_TensorPtr> 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<std::unique_ptr<ResourceHandle>, 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<TF_Tensor*>(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<Safe_TF_TensorPtr> input_vals_safe;
gtl::InlinedVector<std::unique_ptr<ResourceHandle>, 4> resource_handles;
for (PyObject* ndarray : input_ndarrays) {
input_vals_safe.emplace_back(make_safe(static_cast<TF_Tensor*>(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;