experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Allocate input/output buffers at the correct memory space.

Browse Source 
PiperOrigin-RevId: 317761146
Change-Id: Ia27b6765a547451575240e76ea419f705e29bd32
This commit is contained in:
Berkin Ilbeyi 2020-06-22 16:42:36 -07:00 committed by TensorFlower Gardener
parent d2b35a7955
commit 2f9984ff83
2 changed files with 3 additions and 2 deletions
tensorflow/compiler/xla/service
interpreter
executor.cc
transfer_manager.cc

1
tensorflow/compiler/xla/service/interpreter/executor.cc
View File

@ -35,7 +35,6 @@ XlaInterpreterExecutor::~XlaInterpreterExecutor() {}
DeviceMemoryBase XlaInterpreterExecutor::Allocate(uint64 size,
int64 memory_space) {
CHECK_EQ(memory_space, 0);
return DeviceMemoryBase(new char[size], size);
}

4
tensorflow/compiler/xla/service/transfer_manager.cc
View File

@ -355,7 +355,9 @@ StatusOr<ScopedShapedBuffer> TransferManager::AllocateScopedShapedBuffer(
ShapeUtil::GetSubshape(shaped_buffer.on_device_shape(), index);
TF_ASSIGN_OR_RETURN(auto memory,
allocator->Allocate(shaped_buffer.device_ordinal(),
GetByteSizeRequirement(subshape)));
GetByteSizeRequirement(subshape),
/*retry_on_failure=*/true,
subshape.layout().memory_space()));
// Move the allocated buffer into the ScopedShapedBuffer, which owns it.
memory_base = memory.Release();
}