From 13a4593d8b5a6c85ef722365ef4b3b63288bf958 Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Tue, 16 Apr 2019 14:58:41 -0700
Subject: [PATCH] [NFC] Change redzone allocator to use GetSubBuffer instead of
 raw pointer arithmetic

PiperOrigin-RevId: 243885539
---
 .../xla/service/gpu/redzone_allocator.cc      | 25 ++++++++++++-------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/redzone_allocator.cc b/tensorflow/compiler/xla/service/gpu/redzone_allocator.cc
index 30b666448e1..a70ad5a2ee6 100644
--- a/tensorflow/compiler/xla/service/gpu/redzone_allocator.cc
+++ b/tensorflow/compiler/xla/service/gpu/redzone_allocator.cc
@@ -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 {