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:
parent
1333e77450
commit
0b3a25d684
@ -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();
|
||||
|
@ -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."));
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -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;
|
||||
|
Loading…
Reference in New Issue
Block a user