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

@@ -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) {
+  if (dtype != TF_STRING && dtype != TF_RESOURCE &&
-    // Copy the data into a buffer that satisfies Eigen's alignment
+      tensorflow::DataTypeCanUseMemcpy(static_cast<DataType>(dtype)) &&
-    // requirements.
+      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();

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<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."));
   }
 }
 

tensorflow/python/client/tf_session_helper.cc

@@ -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);
+    dims[0] = TF_TensorByteSize(tensor);
-    ResourceHandle* resource_handle =
-        reinterpret_cast<ResourceHandle*>(TF_TensorData(tensor));
-    dims[0] = resource_handle->SerializeAsString().size();
     *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) !=
+  if (TF_TensorType(tensor.get()) == TF_STRING) {
-      static_cast<int64>(TF_TensorByteSize(tensor.get()))) {
+    // Copy element by element.
-    if (TF_TensorType(tensor.get()) == TF_STRING) {
+    auto iter = tensorflow::make_safe(PyArray_IterNew(safe_out_array.get()));
-      // Copy element by element.
+    for (tensorflow::int64 i = 0; i < nelems; ++i) {
-      auto iter = tensorflow::make_safe(PyArray_IterNew(safe_out_array.get()));
+      auto s = CopyStringToPyArrayElement(py_array, iter.get(), tensor.get(),
-      for (tensorflow::int64 i = 0; i < nelems; ++i) {
+                                          nelems, i);
-        auto s = CopyStringToPyArrayElement(py_array, iter.get(), tensor.get(),
+      if (!s.ok()) {
-                                            nelems, i);
+        return s;
-        if (!s.ok()) {
-          return s;
-        }
-        PyArray_ITER_NEXT(iter.get());
       }
-    } else if (TF_TensorType(tensor.get()) == TF_RESOURCE) {
+      PyArray_ITER_NEXT(iter.get());
-      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");
     }
+  } 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
+// `out_tensor` must be non-null. Caller retains ownership of `ndarray`.
-// the deserialized handle. Otherwise it is set to nullptr. Caller becomes owner
+Status PyArrayToTFTensor(PyObject* ndarray, Safe_TF_TensorPtr* out_tensor) {
-// 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) {
   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)),
+    size_t size = PyArray_NBYTES(array);
-                            PyArray_NBYTES(array));
+    array_safe.release();
-    *resource_handle = new ResourceHandle();
+    *out_tensor = make_safe(TF_NewTensor(dtype, {}, 0, PyArray_DATA(array),
-    (*resource_handle)->ParseFromString(serialized);
+                                         size, &DelayedNumpyDecref, array));
-    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);
   } 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());
-    s = PyArrayToTFTensor(value, &inputs_safe.back(), &resource_handle);
     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());
-    s = PyArrayToTFTensor(ndarray, &input_vals_safe.back(), &resource_handle);
-    if (resource_handle != nullptr) {
-      resource_handles.emplace_back(resource_handle);
-    }
     if (!s.ok()) {
       Set_TF_Status_from_Status(out_status, s);
       return;