[NFC] Change redzone allocator to use GetSubBuffer instead of raw pointer arithmetic
PiperOrigin-RevId: 243885539
This commit is contained in:
parent
1bde2c4ce9
commit
13a4593d8b
@ -42,19 +42,27 @@ StatusOr<se::DeviceMemory<uint8>> RedzoneAllocator::AllocateBytes(
|
||||
/*retry_on_failure=*/false));
|
||||
allocated_bytes_excluding_redzones_ += byte_size;
|
||||
|
||||
char* addr =
|
||||
reinterpret_cast<char*>(allocated_buffer.AsDeviceMemoryBase().opaque());
|
||||
se::DeviceMemoryBase lhs_redzone(addr, redzone_size_);
|
||||
static_assert(sizeof(uint8) == 1, "Unexpected size");
|
||||
se::DeviceMemory<uint8> allocated_buffer_memory(
|
||||
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:
|
||||
// - 0 to kRhsRedzoneAlign bytes adjacent to the user buffer, followed by
|
||||
// - redzone_size_ bytes.
|
||||
// We do this because Stream::ThenMemset32 requires the buffer address and
|
||||
// size to be aligned to 4 bytes.
|
||||
se::DeviceMemoryBase rhs_redzone_slop(addr + 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_slop = stream->parent()->GetSubBuffer(
|
||||
&allocated_buffer_memory, redzone_size_ + byte_size, rhs_slop);
|
||||
|
||||
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_,
|
||||
redzone_pattern_};
|
||||
@ -67,8 +75,7 @@ StatusOr<se::DeviceMemory<uint8>> RedzoneAllocator::AllocateBytes(
|
||||
stream->ThenMemset32(&rhs_redzone_nonslop, pattern32, redzone_size_);
|
||||
|
||||
allocated_buffers_.emplace_back(std::move(allocated_buffer), byte_size);
|
||||
return se::DeviceMemory<uint8>(
|
||||
se::DeviceMemoryBase(addr + redzone_size_, byte_size));
|
||||
return data_chunk;
|
||||
}
|
||||
|
||||
Status RedzoneAllocator::CheckRedzones(se::Stream* stream) const {
|
||||
|
Loading…
x
Reference in New Issue
Block a user