Skip to content

Commit

Permalink
Remove pointer arithmetic in StorageObj::AllocNDArray (#7890)
Browse files Browse the repository at this point in the history
The data pointers returned by AllocDataSpace are intended to be opaque
handlers, where previous implementation assumed pointer arithmetic is
valid on them.  Updated to instead use the byte_offset field to
indicate the offset in the allocated array.

Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
  • Loading branch information
Lunderberg and Lunderberg authored Apr 27, 2021
1 parent f681359 commit dac7668
Showing 1 changed file with 2 additions and 7 deletions.
9 changes: 2 additions & 7 deletions src/runtime/vm/memory_manager.cc
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,8 @@ NDArray StorageObj::AllocNDArray(size_t offset, std::vector<int64_t> shape, DLDa

// crtical zone: allocate header, cannot throw
NDArray::Container* container =
new NDArray::Container(nullptr, shape, dtype, this->buffer.device);
new NDArray::Container(this->buffer.data, shape, dtype, this->buffer.device);
container->dl_tensor.byte_offset = offset;

container->SetDeleter(StorageObj::Deleter);
size_t needed_size = GetDataSize(container->dl_tensor);
Expand All @@ -94,12 +95,6 @@ NDArray StorageObj::AllocNDArray(size_t offset, std::vector<int64_t> shape, DLDa
// buffer intact.
container->manager_ctx = reinterpret_cast<void*>(this);

// is this UB?
// The only change we make w.r.t offset is modifying the data pointer
// of the backing tensor to point into the buffer instead of its start.
auto offset_ptr = reinterpret_cast<uint8_t*>(this->buffer.data) + offset;
container->dl_tensor.data = reinterpret_cast<void*>(offset_ptr);

NDArray ret(GetObjectPtr<Object>(container));
// RAII in effect, now run the check.

Expand Down

0 comments on commit dac7668

Please sign in to comment.