Remove a series of allocations and memcpys from the GCS read path (TF infeed path).

The GCS client uses an intermediate std::vector<char> buffer for all reads. This turns out to have a substantial cost, in allocations and in memcpys. During load testing with a simple tool for measuring GCS throughput (but not running TensorFlow), removing this allocation and memcpy increased throughput by about 8%, and significantly reduced steady-state heap memory used.  I have not measured the impact when running the full TensorFlow pipeline, but I expect there to be gains, due to the lower memory utilization and the lower cache demand.

PiperOrigin-RevId: 179581563

1 files changed, 5 insertions, 4 deletions

diff --git a/tensorflow/core/platform/cloud/file_block_cache.h b/tensorflow/core/platform/cloud/file_block_cache.h
index 36dbf9db83..74e792a625 100644
--- a/tensorflow/core/platform/cloud/file_block_cache.h
+++ b/tensorflow/core/platform/cloud/file_block_cache.h
@@ -43,8 +43,9 @@ class FileBlockCache {
   /// cache is constructed. The returned Status should be OK as long as the
   /// read from the remote filesystem succeeded (similar to the semantics of the
   /// read(2) system call).
-  typedef std::function<Status(const string&, size_t, size_t,
-                               std::vector<char>*)>
+  typedef std::function<Status(const string& filename, size_t offset,
+                               size_t buffer_size, char* buffer,
+                               size_t* bytes_transferred)>
       BlockFetcher;
 
   FileBlockCache(size_t block_size, size_t max_bytes, uint64 max_staleness,
@@ -83,8 +84,8 @@ class FileBlockCache {
   ///    placed in `out`.
   /// 4) OK otherwise (i.e. the read succeeded, and at least one byte was placed
   ///    in `out`).
-  Status Read(const string& filename, size_t offset, size_t n,
-              std::vector<char>* out);
+  Status Read(const string& filename, size_t offset, size_t n, char* buffer,
+              size_t* bytes_transferred);
 
   /// Remove all cached blocks for `filename`.
   void RemoveFile(const string& filename) LOCKS_EXCLUDED(mu_);