Skip to content
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

perf: Support read_unsafe in VirtualMemory and RestrictedMemory #256

Merged
merged 1 commit into from
Dec 2, 2024
Merged

perf: Support read_unsafe in VirtualMemory and RestrictedMemory #256

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
14 changes: 7 additions & 7 deletions canbench_results.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ benches:
scopes: {}
btreemap_get_blob_512_1024_v2_mem_manager:
total:
instructions: 3093928908
instructions: 2755579551
heap_increase: 0
stable_memory_increase: 0
scopes: { }
Expand Down Expand Up @@ -139,7 +139,7 @@ benches:
scopes: { }
btreemap_get_u64_u64_v2_mem_manager:
total:
instructions: 619036864
instructions: 522662446
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down Expand Up @@ -223,7 +223,7 @@ benches:
scopes: { }
btreemap_insert_blob_1024_512_v2_mem_manager:
total:
instructions: 6188501297
instructions: 5594888576
heap_increase: 0
stable_memory_increase: 256
scopes: {}
Expand Down Expand Up @@ -379,7 +379,7 @@ benches:
scopes: {}
btreemap_insert_u64_u64_mem_manager:
total:
instructions: 937850849
instructions: 830923183
heap_increase: 0
stable_memory_increase: 0
scopes: { }
Expand Down Expand Up @@ -631,7 +631,7 @@ benches:
scopes: {}
memory_manager_overhead:
total:
instructions: 1182056741
instructions: 1182056386
heap_increase: 0
stable_memory_increase: 8320
scopes: {}
Expand Down Expand Up @@ -661,7 +661,7 @@ benches:
scopes: {}
vec_get_blob_4_mem_manager:
total:
instructions: 12856122
instructions: 12194073
heap_increase: 0
stable_memory_increase: 0
scopes: { }
Expand All @@ -673,7 +673,7 @@ benches:
scopes: { }
vec_get_blob_64_mem_manager:
total:
instructions: 22526389
instructions: 20656466
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down
5 changes: 5 additions & 0 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,11 @@ impl<M: Memory> Memory for RestrictedMemory<M> {
.read(self.page_range.start * WASM_PAGE_SIZE + offset, dst)
}

unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) {
self.memory
.read_unsafe(self.page_range.start * WASM_PAGE_SIZE + offset, dst, count)
}

fn write(&self, offset: u64, src: &[u8]) {
self.memory
.write(self.page_range.start * WASM_PAGE_SIZE + offset, src)
Expand Down
86 changes: 61 additions & 25 deletions src/memory_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,16 +171,16 @@ struct Header {

version: u8,

// The number of buckets allocated by the memory manager.
/// The number of buckets allocated by the memory manager.
num_allocated_buckets: u16,

// The size of a bucket in Wasm pages.
/// The size of a bucket in Wasm pages.
bucket_size_in_pages: u16,

// Reserved bytes for future extensions
/// Reserved bytes for future extensions
_reserved: [u8; HEADER_RESERVED_BYTES],

// The size of each individual memory that can be created by the memory manager.
/// The size of each individual memory that can be created by the memory manager.
memory_sizes_in_pages: [u64; MAX_NUM_MEMORIES as usize],
}

Expand Down Expand Up @@ -209,6 +209,12 @@ impl<M: Memory> Memory for VirtualMemory<M> {
self.memory_manager.borrow().read(self.id, offset, dst)
}

unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) {
self.memory_manager
.borrow()
.read_unsafe(self.id, offset, dst, count)
}

fn write(&self, offset: u64, src: &[u8]) {
self.memory_manager.borrow().write(self.id, offset, src)
}
Expand All @@ -218,15 +224,15 @@ impl<M: Memory> Memory for VirtualMemory<M> {
struct MemoryManagerInner<M: Memory> {
memory: M,

// The number of buckets that have been allocated.
/// The number of buckets that have been allocated.
allocated_buckets: u16,

bucket_size_in_pages: u16,

// An array storing the size (in pages) of each of the managed memories.
/// An array storing the size (in pages) of each of the managed memories.
memory_sizes_in_pages: [u64; MAX_NUM_MEMORIES as usize],

// A map mapping each managed memory to the bucket ids that are allocated to it.
/// A map mapping each managed memory to the bucket ids that are allocated to it.
memory_buckets: Vec<Vec<BucketId>>,
}

Expand Down Expand Up @@ -313,12 +319,12 @@ impl<M: Memory> MemoryManagerInner<M> {
write_struct(&header, Address::from(0), &self.memory);
}

// Returns the size of a memory (in pages).
/// Returns the size of a memory (in pages).
fn memory_size(&self, id: MemoryId) -> u64 {
self.memory_sizes_in_pages[id.0 as usize]
}

// Grows the memory with the given id by the given number of pages.
/// Grows the memory with the given id by the given number of pages.
fn grow(&mut self, id: MemoryId, pages: u64) -> i64 {
// Compute how many additional buckets are needed.
let old_size = self.memory_size(id);
Expand Down Expand Up @@ -384,23 +390,35 @@ impl<M: Memory> MemoryManagerInner<M> {
}
}

#[inline]
fn read(&self, id: MemoryId, offset: u64, dst: &mut [u8]) {
if (offset + dst.len() as u64) > self.memory_size(id) * WASM_PAGE_SIZE {
// SAFETY: this is trivially safe because dst has dst.len() space.
unsafe { self.read_unsafe(id, offset, dst.as_mut_ptr(), dst.len()) }
}

/// # Safety
///
/// Callers must guarantee that
/// * it is valid to write `count` number of bytes starting from `dst`,
/// * `dst..dst + count` does not overlap with `self`.
unsafe fn read_unsafe(&self, id: MemoryId, offset: u64, dst: *mut u8, count: usize) {
if (offset + count as u64) > self.memory_size(id) * WASM_PAGE_SIZE {
panic!("{id:?}: read out of bounds");
}

let mut bytes_read = 0;
for Segment { address, length } in self.bucket_iter(id, offset, dst.len()) {
self.memory.read(
address.get(),
&mut dst[bytes_read as usize..(bytes_read + length.get()) as usize],
);
let mut bytes_read: usize = 0;
for Segment { address, length } in self.bucket_iter(id, offset, count) {
let length = length.get().try_into().expect("Length overflows usize");
self.memory
.read_unsafe(address.get(), dst.add(bytes_read), length);

bytes_read += length.get();
bytes_read = bytes_read
.checked_add(length)
.expect("Bytes read overflowed usize");
}
}

// Initializes a [`BucketIterator`].
/// Initializes a [`BucketIterator`].
fn bucket_iter(&self, MemoryId(id): MemoryId, offset: u64, length: usize) -> BucketIterator {
// Get the buckets allocated to the given memory id.
let buckets = self.memory_buckets[id as usize].as_slice();
Expand All @@ -419,13 +437,13 @@ impl<M: Memory> MemoryManagerInner<M> {
Bytes::from(self.bucket_size_in_pages as u64 * WASM_PAGE_SIZE)
}

// Returns the number of buckets needed to accommodate the given number of pages.
/// Returns the number of buckets needed to accommodate the given number of pages.
fn num_buckets_needed(&self, num_pages: u64) -> u64 {
// Ceiling division.
(num_pages + self.bucket_size_in_pages as u64 - 1) / self.bucket_size_in_pages as u64
}

// Returns the underlying memory.
/// Returns the underlying memory.
pub fn into_memory(self) -> M {
self.memory
}
Expand Down Expand Up @@ -509,7 +527,7 @@ impl Iterator for BucketIterator<'_> {
}

impl<'a> BucketIterator<'a> {
// Returns the address of a given bucket.
/// Returns the address of a given bucket.
fn bucket_address(&self, id: BucketId) -> Address {
Address::from(BUCKETS_OFFSET_IN_BYTES) + self.bucket_size_in_bytes * Bytes::from(id.0)
}
Expand Down Expand Up @@ -815,10 +833,28 @@ mod test {
// Write a random blob into the memory, growing the memory as it needs to.
write(memory, offset, &data);

// Verify the blob can be read back.
let mut bytes = vec![0; data.len()];
memory.read(offset, &mut bytes);
assert_eq!(bytes, data);
{
// Verify the blob can be read back using read.
let mut bytes = vec![0; data.len()];
memory.read(offset, &mut bytes);
assert_eq!(bytes, data);
}

{
// Verify the blob can be read back using read_to_vec.
let mut bytes = vec![];
read_to_vec(memory, offset.into(), &mut bytes, data.len());
assert_eq!(bytes, data);
}

{
// Verify the blob can be read back using read_unsafe.
let mut bytes = vec![0; data.len()];
unsafe {
memory.read_unsafe(offset, bytes.as_mut_ptr(), data.len());
}
assert_eq!(bytes, data);
}
}
});
}
Expand Down