diff options
author | 2017-11-14 13:48:43 -0800 | |
---|---|---|
committer | 2017-11-14 13:52:47 -0800 | |
commit | b77b20f8e8ef5a670201032a8fc8daf157524b74 (patch) | |
tree | adc6745b585b4f2f4410422de25ef96c6ab96cc7 | |
parent | 29d84f18369cfe08beae97cff0aa8bde601b4cfc (diff) |
Automated g4 rollback of changelist 175571632
PiperOrigin-RevId: 175729221
-rw-r--r-- | tensorflow/compiler/xla/service/shaped_buffer.cc | 8 | ||||
-rw-r--r-- | tensorflow/compiler/xla/service/shaped_buffer.h | 4 |
2 files changed, 12 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/shaped_buffer.cc b/tensorflow/compiler/xla/service/shaped_buffer.cc index a2a442eb1a..a57ebf59e7 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.cc +++ b/tensorflow/compiler/xla/service/shaped_buffer.cc @@ -63,6 +63,14 @@ void ShapedBuffer::clear() { } } +void ShapedBuffer::AddBufferAtIndex( + const perftools::gputools::DeviceMemoryBase& buffer, + const ShapeIndex& shape_index) { + *mutable_shape_index_to_buffer_entry()->mutable_element(shape_index) = + buffers().size(); + mutable_buffers()->push_back(buffer); +} + const se::DeviceMemoryBase& ShapedBuffer::buffer( const ShapeIndex& index) const { return buffers_[shape_index_to_buffer_entry_.element(index)]; diff --git a/tensorflow/compiler/xla/service/shaped_buffer.h b/tensorflow/compiler/xla/service/shaped_buffer.h index e5ea06fb13..b440948700 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.h +++ b/tensorflow/compiler/xla/service/shaped_buffer.h @@ -75,6 +75,10 @@ class ShapedBuffer { // Set all device memory pointers in the object to null. void clear(); + // Adds a new buffer at the given shape index. + void AddBufferAtIndex(const perftools::gputools::DeviceMemoryBase& buffer, + const ShapeIndex& shape_index); + protected: // The shape of the device buffer with layout. const Shape shape_; |