Skip to content

perf: Add new efficient API read_to_vec #246

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

perf: Add new efficient API read_to_vec #246

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
212 changes: 106 additions & 106 deletions canbench_results.yml
View file
Open in desktop

Large diffs are not rendered by default.

5 changes: 2 additions & 3 deletions src/base_vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,11 +245,10 @@ impl<T: Storable, M: Memory> BaseVec<T, M> {
}

/// Reads the item at the specified index without any bound checks.
fn read_entry_to(&self, index: u64, buf: &mut std::vec::Vec<u8>) {
fn read_entry_to(&self, index: u64, buf: &mut Vec<u8>) {
let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
let (data_offset, data_size) = self.read_entry_size(offset);
buf.resize(data_size, 0);
self.memory.read(data_offset, &mut buf[..]);
self.memory.read_to_vec(data_offset, data_size, buf);
}

/// Sets the vector's length.
Expand Down
4 changes: 2 additions & 2 deletions src/btreemap/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,8 @@ impl<K: Storable + Ord + Clone> Node<K> {
};

let value_len = read_u32(&reader, Address::from(offset.get())) as usize;
let mut bytes = vec![0; value_len];
reader.read((offset + U32_SIZE).get(), &mut bytes);
let mut bytes = vec![];
reader.read_to_vec((offset + U32_SIZE).get(), value_len, &mut bytes);

bytes
}
Expand Down
12 changes: 12 additions & 0 deletions src/btreemap/node/io.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,18 @@ impl<'a, M: Memory> Memory for NodeReader<'a, M> {
}
}

fn read_to_vec(&self, offset: u64, len: usize, dst: &mut Vec<u8>) {
// If the read is only in the initial page, then read it directly in one go.
// This is a performance enhancement.
if (offset + len as u64) < self.page_size.get() as u64 {
self.memory
.read_to_vec(self.address.get() + offset, len, dst);
} else {
dst.resize(len, 0);
self.read(offset, &mut dst[..]);
}
}

fn write(&self, _: u64, _: &[u8]) {
unreachable!("NodeReader does not support write")
}
Expand Down
5 changes: 2 additions & 3 deletions src/btreemap/node/v1.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,15 +65,14 @@ impl<K: Storable + Ord + Clone> Node<K> {
// Load the entries.
let mut keys_encoded_values = Vec::with_capacity(header.num_entries as usize);
let mut offset = NodeHeader::size();
let mut buf = Vec::with_capacity(max_key_size.max(max_value_size) as usize);
let mut buf = vec![];
for _ in 0..header.num_entries {
// Read the key's size.
let key_size = read_u32(memory, address + offset);
offset += U32_SIZE;

// Read the key.
buf.resize(key_size as usize, 0);
memory.read((address + offset).get(), &mut buf);
memory.read_to_vec((address + offset).get(), key_size as usize, &mut buf);
offset += Bytes::from(max_key_size);
let key = K::from_bytes(Cow::Borrowed(&buf));
// Values are loaded lazily. Store a reference and skip loading it.
Expand Down
3 changes: 1 addition & 2 deletions src/btreemap/node/v2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,8 +164,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
};

// Load the key.
buf.resize(key_size as usize, 0);
reader.read(offset.get(), &mut buf);
reader.read_to_vec(offset.get(), key_size as usize, &mut buf);
let key = K::from_bytes(Cow::Borrowed(&buf));
offset += Bytes::from(key_size);
keys_encoded_values.push((key, Value::by_ref(Bytes::from(0usize))));
Expand Down
4 changes: 2 additions & 2 deletions src/cell.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,8 @@ impl<T: Storable, M: Memory> Cell<T, M> {
///
/// PRECONDITION: memory is large enough to contain the value.
fn read_value(memory: &M, len: u32) -> T {
let mut buf = vec![0; len as usize];
memory.read(HEADER_V1_SIZE, &mut buf);
let mut buf = vec![];
memory.read_to_vec(HEADER_V1_SIZE, len as usize, &mut buf);
T::from_bytes(Cow::Owned(buf))
}

Expand Down
12 changes: 12 additions & 0 deletions src/ic0_memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,18 @@ impl Memory for Ic0StableMemory {
unsafe { stable64_read(dst.as_ptr() as u64, offset, dst.len() as u64) }
}

fn read_to_vec(&self, offset: u64, len: usize, dst: &mut Vec<u8>) {
dst.clear();
dst.reserve(len);
unsafe {
// SAFETY: This is safe because of the ic0 api guarantees.
stable64_read(dst.as_ptr() as u64, offset, len as u64);
// SAFETY: This is safe because stable64_read guarantees that the first len bytes
// will be initialized.
dst.set_len(len);
}
}

fn write(&self, offset: u64, src: &[u8]) {
// SAFETY: This is safe because of the ic0 api guarantees.
unsafe { stable64_write(offset, src.as_ptr() as u64, src.len() as u64) }
Expand Down
10 changes: 10 additions & 0 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,16 @@ pub trait Memory {
/// and replaces the corresponding bytes in dst.
fn read(&self, offset: u64, dst: &mut [u8]);

/// Copies len bytes of data starting from offset out of the stable memory into dst.
/// Implementations must not make any assumptions about dst (e.g. length, contents, capacity).
/// Callers are allowed to pass empty vectors. After the method returns, dst.len() == len.
/// This method is an alternative to read which does not require initializing a buffer and may
/// therefore be faster.
fn read_to_vec(&self, offset: u64, len: usize, dst: &mut std::vec::Vec<u8>) {
dst.resize(len, 0);
self.read(offset, &mut dst[..]);
}
Comment on lines +63 to +66
Copy link
Member

@dsarlis dsarlis Nov 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure if we want to have this be part of the Memory trait. It depends on how religious we want to be here.

The Memory trait is supposed to model the operations that are available for stable memory which in turn models the operations available for heap memory in Wasm. There's no read_to_vec so following that logic, maybe we avoid adding it to the trait.

We can probably have it as a free function or you think there's a significant advantage in having it on the trait?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The benefit of having it on the trait is that implementors of the trait can decide to override it in an efficient way. For example, Ic0StableMemory can read directly read into a Vec skipping initialization.

is supposed to model the operations that are available for stable memory

That's a valid point. Another alternative could be changing the read method to read into dst: &mut [MaybeUninit<u8>] which is technically the truth. It's unfortunately an unsolved problem in Rust to create uninitialized [u8] slices (something totally acceptable in C or C++). The documentation is very explicit about warning people not to create slices with uninitialized values even if those values will never be read.

Copy link
Contributor Author

@frankdavid frankdavid Nov 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more alternative would be changing the read method to take a raw pointer and a length. In such cases, Rust does not require any initialization. Callers could then do:

let mut v = Vec::with_capacity(len);
memory.read(offset, len, v.as_mut_ptr());
v.set_len(x);

However, read would likely need to be marked as unsafe since clients would need to guarantee that len number of bytes after the passed pointer are writable.

Copy link
Contributor Author

@frankdavid frankdavid Nov 21, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tested what if we replace read_to_vec with unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize)

It allows for even better optimizations than read_to_vec leading to up to 40% instruction reduction. This function is unsafe because clients must promise that count number of bytes after dst are writeable.

(We could still keep read_to_vec but make it into an associated method like you mentioned. It would provide a safe and efficient way for clients to read and would use read_unsafe under the hood.)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, to make sure I understand your suggestion correctly, it would be:

  1. Add unsafe fn read_unsafe for maximum performance if clients know what they are doing -- not super excited because I think it can be a big trap but as long as we provide safe alternatives I'm not strongly opposed.
  2. Add read_to_vec as an associated method, not on the trait.
  3. We would also leave the original read but it's the slowest of them all?

Thanks for all the exploration around possible options.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's correct. I'd keep the original read because sometimes you have a fixed-size array that you want to fill. If the array is small, initializing it is fast and using a vector is unnecessary (and slower). read is the best option for this use-case.

When reading a larger blob, it's better not to allocate on the stack, so reading into a Vec is the best approach. read_to_vec can be used for that.

And Memory implementations must be able to read into uninitialized buffers (pointer + length) which is trivial for the mostly used Ic0Memory implementation. The default implementation will just initialize the buffer with 0-s first and then call the existing read(). If we make read_to_vec into an associated function, we don't have to pollute the Memory interface with Vec.

if clients know what they are doing

The API will be the same as for example std::ptr::copy with the same requirements. And we can provide safe wrappers while staying efficient.

I'll send a new PR with this new approach.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bit of a random question, but we have several canisters that use stable-structures, and we'd love to benefit from these improvements - would we be able to just by upgrading the crate once this change is out?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@rikonor Yeah, we'll need to cut a new release because there's a few optimizations already included that are unreleased. Maybe we can target to make a release after this change with read_to_vec is merged.


/// Copies the data referred to by src and replaces the
/// corresponding segment starting at offset in the stable memory.
fn write(&self, offset: u64, src: &[u8]);
Expand Down
4 changes: 2 additions & 2 deletions src/log.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,8 +331,8 @@ impl<T: Storable, INDEX: Memory, DATA: Memory> Log<T, INDEX, DATA> {
/// ignores the result.
pub fn read_entry(&self, idx: u64, buf: &mut Vec<u8>) -> Result<(), NoSuchEntry> {
let (offset, len) = self.entry_meta(idx).ok_or(NoSuchEntry)?;
buf.resize(len, 0);
self.data_memory.read(HEADER_OFFSET + offset, buf);
self.data_memory
.read_to_vec(HEADER_OFFSET + offset, len, buf);
Ok(())
}

Expand Down
12 changes: 8 additions & 4 deletions src/memory_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,9 +239,9 @@ impl<M: Memory> MemoryManagerInner<M> {
}

// Check if the magic in the memory corresponds to this object.
let mut dst = vec![0; 3];
let mut dst = [0; 3];
memory.read(0, &mut dst);
if dst != MAGIC {
if &dst != MAGIC {
// No memory manager found. Create a new instance.
MemoryManagerInner::new(memory, bucket_size_in_pages)
} else {
Expand Down Expand Up @@ -277,8 +277,12 @@ impl<M: Memory> MemoryManagerInner<M> {
assert_eq!(&header.magic, MAGIC, "Bad magic.");
assert_eq!(header.version, LAYOUT_VERSION, "Unsupported version.");

let mut buckets = vec![0; MAX_NUM_BUCKETS as usize];
memory.read(bucket_allocations_address(BucketId(0)).get(), &mut buckets);
let mut buckets = vec![];
memory.read_to_vec(
bucket_allocations_address(BucketId(0)).get(),
MAX_NUM_BUCKETS as usize,
&mut buckets,
);

let mut memory_buckets = BTreeMap::new();
for (bucket_idx, memory) in buckets.into_iter().enumerate() {
Expand Down
13 changes: 13 additions & 0 deletions src/vec_mem.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,16 @@ impl Memory for RefCell<Vec<u8>> {
dst.copy_from_slice(&self.borrow()[offset as usize..n as usize]);
}

fn read_to_vec(&self, offset: u64, len: usize, dst: &mut Vec<u8>) {
let n = offset.checked_add(len as u64).expect("read: out of bounds");

if n as usize > self.borrow().len() {
panic!("read: out of bounds");
}
dst.clear();
dst.extend_from_slice(&self.borrow()[offset as usize..n as usize]);
}

fn write(&self, offset: u64, src: &[u8]) {
let n = offset
.checked_add(src.len() as u64)
Expand All @@ -61,6 +71,9 @@ impl<M: Memory> Memory for Rc<M> {
fn read(&self, offset: u64, dst: &mut [u8]) {
self.deref().read(offset, dst)
}
fn read_to_vec(&self, offset: u64, len: usize, dst: &mut Vec<u8>) {
self.deref().read_to_vec(offset, len, dst)
}
fn write(&self, offset: u64, src: &[u8]) {
self.deref().write(offset, src)
}
Expand Down