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

[NFC] Change redzone allocator to use GetSubBuffer instead of raw pointer arithmetic

Browse Source 
PiperOrigin-RevId: 243885539
This commit is contained in:
George Karpenkov 2019-04-16 14:58:41 -07:00 committed by TensorFlower Gardener
parent 1bde2c4ce9
commit 13a4593d8b
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
tensorflow/compiler/xla/service/gpu/redzone_allocator.cc
View File

@ -42,19 +42,27 @@ StatusOr<se::DeviceMemory<uint8>> RedzoneAllocator::AllocateBytes(
/*retry_on_failure=*/false)); /*retry_on_failure=*/false));
allocated_bytes_excluding_redzones_ += byte_size; allocated_bytes_excluding_redzones_ += byte_size;
char* addr = static_assert(sizeof(uint8) == 1, "Unexpected size");
reinterpret_cast<char*>(allocated_buffer.AsDeviceMemoryBase().opaque()); se::DeviceMemory<uint8> allocated_buffer_memory(
se::DeviceMemoryBase lhs_redzone(addr, redzone_size_); allocated_buffer.AsDeviceMemoryBase());
se::DeviceMemory<uint8> lhs_redzone = stream->parent()->GetSubBuffer(
&allocated_buffer_memory, 0, redzone_size_);
se::DeviceMemory<uint8> data_chunk = stream->parent()->GetSubBuffer(
&allocated_buffer_memory, redzone_size_, byte_size);
// Split up the RHS redzone into two pieces: // Split up the RHS redzone into two pieces:
// - 0 to kRhsRedzoneAlign bytes adjacent to the user buffer, followed by // - 0 to kRhsRedzoneAlign bytes adjacent to the user buffer, followed by
// - redzone_size_ bytes. // - redzone_size_ bytes.
// We do this because Stream::ThenMemset32 requires the buffer address and // We do this because Stream::ThenMemset32 requires the buffer address and
// size to be aligned to 4 bytes. // size to be aligned to 4 bytes.
se::DeviceMemoryBase rhs_redzone_slop(addr + redzone_size_ + byte_size, se::DeviceMemory<uint8> rhs_redzone_slop = stream->parent()->GetSubBuffer(
rhs_slop); &allocated_buffer_memory, redzone_size_ + byte_size, rhs_slop);
se::DeviceMemoryBase rhs_redzone_nonslop(
addr + redzone_size_ + byte_size + rhs_slop, redzone_size_); se::DeviceMemory<uint8> rhs_redzone_nonslop = stream->parent()->GetSubBuffer(
&allocated_buffer_memory, redzone_size_ + byte_size + rhs_slop,
redzone_size_);
uint8 pattern_arr[] = {redzone_pattern_, redzone_pattern_, redzone_pattern_, uint8 pattern_arr[] = {redzone_pattern_, redzone_pattern_, redzone_pattern_,
redzone_pattern_}; redzone_pattern_};
@ -67,8 +75,7 @@ StatusOr<se::DeviceMemory<uint8>> RedzoneAllocator::AllocateBytes(
stream->ThenMemset32(&rhs_redzone_nonslop, pattern32, redzone_size_); stream->ThenMemset32(&rhs_redzone_nonslop, pattern32, redzone_size_);
allocated_buffers_.emplace_back(std::move(allocated_buffer), byte_size); allocated_buffers_.emplace_back(std::move(allocated_buffer), byte_size);
return se::DeviceMemory<uint8>( return data_chunk;
se::DeviceMemoryBase(addr + redzone_size_, byte_size));
} }
Status RedzoneAllocator::CheckRedzones(se::Stream* stream) const { Status RedzoneAllocator::CheckRedzones(se::Stream* stream) const {