diff options
| -rw-r--r-- | tensorflow/compiler/xla/service/allocation_tracker.cc | 4 | ||||
| -rw-r--r-- | tensorflow/compiler/xla/service/transfer_manager.h | 7 | ||||
| -rw-r--r-- | tensorflow/compiler/xla/tests/hlo_test_base.cc | 4 |
3 files changed, 2 insertions, 13 deletions
diff --git a/tensorflow/compiler/xla/service/allocation_tracker.cc b/tensorflow/compiler/xla/service/allocation_tracker.cc index e59fad4e05..49a621810e 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.cc +++ b/tensorflow/compiler/xla/service/allocation_tracker.cc @@ -125,9 +125,7 @@ tensorflow::Status AllocationTracker::DeallocateShape( handle_map.erase(device_memory->opaque()); } - // TODO(b/36256956) Ideally tuple elements could always be distinct buffers. - if (ShapeUtil::IsTuple(shape) && - backend->transfer_manager()->TupleElementsAreDistinctBuffers()) { + if (ShapeUtil::IsTuple(shape)) { // Traverse into tuple recursively deallocating buffers. TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor, backend->stream_executor(device_ordinal)); diff --git a/tensorflow/compiler/xla/service/transfer_manager.h b/tensorflow/compiler/xla/service/transfer_manager.h index b052bb8146..83e893a14a 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.h +++ b/tensorflow/compiler/xla/service/transfer_manager.h @@ -99,13 +99,6 @@ class TransferManager { // region for a host-to-device transfer. virtual int64 GetByteSizeRequirement(const Shape& shape) = 0; - // Returns whether tuple elements are distinct buffers (in which case each of - // the elements of a tuple should be deallocated, in addition to the tuple's - // buffer itself). - // - // TODO(b/36256956) Ideally tuple elements could always be distinct buffers. - virtual bool TupleElementsAreDistinctBuffers() const { return true; } - // Transfer a memory block of the given size from the device source into the // 'destination' buffer. // diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index 62878fed55..4e956bc00c 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -123,9 +123,7 @@ StatusOr<se::DeviceMemoryBase> HloTestBase::Execute( *result_shape = executable->result_shape(); - // TODO(b/36256956) Ideally tuple elements could always be distinct buffers. - if (ShapeUtil::IsTuple(*result_shape) && - backend_->transfer_manager()->TupleElementsAreDistinctBuffers()) { + if (ShapeUtil::IsTuple(*result_shape)) { // We must record element buffers of tuples as well to avoid leaks. DCHECK(!ShapeUtil::IsNestedTuple(*result_shape)); TF_ASSIGN_OR_RETURN( |