C API: TF_Tensors will always be in host memory.

This change undoes some experimentation in
commit 22651083406ca01ac9d481e3367a3510d25f88cd
and restores TF_Tensor behavior to what is was prior to that change.

PiperOrigin-RevId: 164146670

3 files changed, 18 insertions, 94 deletions

diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc
index d90b653fef..cb074cb2f1 100644
--- a/tensorflow/c/c_api.cc
+++ b/tensorflow/c/c_api.cc
@@ -182,25 +182,7 @@ Status MessageToBuffer(const tensorflow::protobuf::Message& in,
 
 }  // namespace
 
-TF_BufferAndDevice::TF_BufferAndDevice(TensorBuffer* buffer)
-    : buffer_(buffer), device_owner_(nullptr), device_index_(-1) {}
-
-TF_BufferAndDevice::TF_BufferAndDevice(TensorBuffer* buffer,
-                                       TF_Session* session, int device_index)
-    : buffer_(buffer), device_owner_(session), device_index_(device_index) {
-  mutex_lock l(device_owner_->mu);
-  device_owner_->num_outstanding_buffers++;
-}
-
-TF_BufferAndDevice::~TF_BufferAndDevice() {
-  buffer_->Unref();
-  if (device_owner_ != nullptr) {
-    mutex_lock l(device_owner_->mu);
-    device_owner_->num_outstanding_buffers--;
-  }
-}
-
-TF_Tensor::~TF_Tensor() { delete buffer; }
+TF_Tensor::~TF_Tensor() { buffer->Unref(); }
 
 TF_Tensor* TF_AllocateTensor(TF_DataType dtype, const int64_t* dims,
                              int num_dims, size_t len) {
@@ -241,14 +223,14 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims,
     buf->deallocator_ = deallocator;
     buf->deallocator_arg_ = deallocator_arg;
   }
-  return new TF_Tensor{dtype, TensorShape(dimvec), new TF_BufferAndDevice(buf)};
+  return new TF_Tensor{dtype, TensorShape(dimvec), buf};
 }
 
 TF_Tensor* TF_TensorMaybeMove(TF_Tensor* tensor) {
   // It is safe to move the Tensor if and only if we own the unique reference to
   // it. In that case, we might as well not delete and reallocate, but a future
   // implementation might need to do so.
-  TensorBuffer* buf = tensor->buffer->buffer();
+  TensorBuffer* buf = tensor->buffer;
   if (buf->RefCountIsOne() && buf->root_buffer()->RefCountIsOne() &&
       buf->OwnsMemory()) {
     return tensor;
@@ -263,13 +245,8 @@ int TF_NumDims(const TF_Tensor* t) { return t->shape.dims(); }
 int64_t TF_Dim(const TF_Tensor* t, int dim_index) {
   return static_cast<int64_t>(t->shape.dim_size(dim_index));
 }
-size_t TF_TensorByteSize(const TF_Tensor* t) {
-  return t->buffer->buffer()->size();
-}
-void* TF_TensorData(const TF_Tensor* t) {
-  if (t->buffer->on_cpu()) return t->buffer->buffer()->data();
-  return nullptr;
-}
+size_t TF_TensorByteSize(const TF_Tensor* t) { return t->buffer->size(); }
+void* TF_TensorData(const TF_Tensor* t) { return t->buffer->data(); }
 
 // --------------------------------------------------------------------------
 size_t TF_StringEncode(const char* src, size_t src_len, char* dst,
@@ -428,10 +405,6 @@ 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 "
@@ -448,11 +421,7 @@ Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst) {
     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());
+    *dst = TensorCApi::MakeTensor(src->dtype, src->shape, src->buffer);
     return Status::OK();
   }
   // TF_STRING tensors require copying since Tensor class expects a sequence of
@@ -514,7 +483,7 @@ TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src,
     TensorBuffer* buf = TensorCApi::Buffer(src);
     buf->Ref();
     return new TF_Tensor{static_cast<TF_DataType>(src.dtype()), src.shape(),
-                         new TF_BufferAndDevice(buf)};
+                         buf};
   }
   // DT_STRING tensors require a copying since TF_Tensor.buffer expects a flatly
   // encoded sequence of strings.
@@ -2206,11 +2175,7 @@ void TF_AddGradients(TF_Graph* g, TF_Output* y, int ny, TF_Output* x, int nx,
 // TF_Session functions ----------------------------------------------
 
 TF_Session::TF_Session(tensorflow::Session* s, TF_Graph* g)
-    : session(s),
-      graph(g),
-      last_num_graph_nodes(0),
-      device_mgr(nullptr),
-      num_outstanding_buffers(0) {
+    : session(s), graph(g), last_num_graph_nodes(0), device_mgr(nullptr) {
   if (s->LocalDeviceManager(&device_mgr).ok()) {
     devices = device_mgr->ListDevices();
   }
@@ -2299,30 +2264,16 @@ void TF_CloseSession(TF_Session* s, TF_Status* status) {
 }
 
 void TF_DeleteSession(TF_Session* s, TF_Status* status) {
-  {
-    mutex_lock l(s->mu);
-    if (s->num_outstanding_buffers > 0) {
-      // This can probably be relaxed: An alternative might be to mark
-      // this session for deletion and do the actual delete only when
-      // the last TF_BufferAndDevice has been deleted.
-      status->status = FailedPrecondition(
-          s->num_outstanding_buffers,
-          " TF_Tensor objects with memory backed by a device "
-          "owned by this TF_Session are still alive. Release "
-          "them using TF_DeleteTensor and retry");
-      return;
-    }
-    status->status = Status::OK();
-    TF_Graph* const graph = s->graph;
-    if (graph != nullptr) {
-      graph->mu.lock();
-      graph->num_sessions -= 1;
-      const bool del = graph->delete_requested && graph->num_sessions == 0;
-      graph->mu.unlock();
-      if (del) delete graph;
-    }
-    delete s->session;
+  status->status = Status::OK();
+  TF_Graph* const graph = s->graph;
+  if (graph != nullptr) {
+    graph->mu.lock();
+    graph->num_sessions -= 1;
+    const bool del = graph->delete_requested && graph->num_sessions == 0;
+    graph->mu.unlock();
+    if (del) delete graph;
   }
+  delete s->session;
   delete s;
 }
 
diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h
index df38ae5399..43b5078013 100644
--- a/tensorflow/c/c_api.h
+++ b/tensorflow/c/c_api.h
@@ -263,9 +263,6 @@ TF_CAPI_EXPORT extern int64_t TF_Dim(const TF_Tensor* tensor, int dim_index);
 TF_CAPI_EXPORT extern size_t TF_TensorByteSize(const TF_Tensor*);
 
 // Return a pointer to the underlying data buffer.
-//
-// Returns NULL if the underlying data is not in host memory
-// (for example, if it refers to addresses in GPU memory).
 TF_CAPI_EXPORT extern void* TF_TensorData(const TF_Tensor*);
 
 // --------------------------------------------------------------------------
diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h
index 687e18aace..b89acbcf35 100644
--- a/tensorflow/c/c_api_internal.h
+++ b/tensorflow/c/c_api_internal.h
@@ -36,7 +36,6 @@ namespace tensorflow {
 class Device;
 class DeviceMgr;
 }  // namespace tensorflow
-class TF_BufferAndDevice;
 
 // Internal structures used by the C API. These are likely to change and should
 // not be depended on.
@@ -50,7 +49,7 @@ struct TF_Tensor {
 
   TF_DataType dtype;
   tensorflow::TensorShape shape;
-  TF_BufferAndDevice* buffer;
+  tensorflow::TensorBuffer* buffer;
 };
 
 struct TF_SessionOptions {
@@ -120,7 +119,6 @@ struct TF_Session {
   // buffers of a TF_Tensor pinned in device memory.
   const tensorflow::DeviceMgr* device_mgr;   // Owned by session.
   std::vector<tensorflow::Device*> devices;  // Owned by device_mgr.
-  int num_outstanding_buffers GUARDED_BY(mu);
 };
 
 struct TF_ImportGraphDefOptions {
@@ -131,28 +129,6 @@ struct TF_DeviceList {
   std::vector<tensorflow::DeviceAttributes> response;
 };
 
-// TF_BufferAndDevice encapsulates the memory addresses of data backing a Tensor
-// and the device (e.g., GPU or host) whose memory the addresses refer to.
-class TF_BufferAndDevice {
- public:
-  explicit TF_BufferAndDevice(tensorflow::TensorBuffer* buffer);
-  TF_BufferAndDevice(tensorflow::TensorBuffer* buffer, TF_Session* session,
-                     int device_index);
-  ~TF_BufferAndDevice();
-
-  tensorflow::TensorBuffer* buffer() const { return buffer_; }
-  tensorflow::Device* device() const {
-    if (device_owner_ == nullptr) return nullptr;
-    return device_owner_->devices[device_index_];
-  }
-  bool on_cpu() const { return device() == nullptr; }
-
- private:
-  tensorflow::TensorBuffer* buffer_;
-  TF_Session* device_owner_;
-  const int device_index_;
-};
-
 namespace tensorflow {
 
 class TensorCApi {