Skip to content

Commit

Permalink
Optimize buckets
Browse files Browse the repository at this point in the history
  • Loading branch information
@frankdavid
frankdavid committed Dec 2, 2024
1 parent 3b995f5 commit 257abad
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 93 deletions.
48 changes: 24 additions & 24 deletions canbench_results.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,10 @@ benches:
scopes: {}
btreemap_get_blob_512_1024_v2_mem_manager:
total:
instructions: 2755579551
instructions: 2621506893
heap_increase: 0
stable_memory_increase: 0
scopes: { }
scopes: {}
btreemap_get_blob_64_1024:
total:
instructions: 572513751
Expand All @@ -103,13 +103,13 @@ benches:
scopes: {}
btreemap_get_blob_8_u64:
total:
instructions: 196934591
instructions: 196964591
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_get_blob_8_u64_v2:
total:
instructions: 291577086
instructions: 291607086
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand All @@ -127,19 +127,19 @@ benches:
scopes: {}
btreemap_get_u64_u64:
total:
instructions: 177269159
instructions: 177299159
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_get_u64_u64_v2:
total:
instructions: 255485892
instructions: 255515892
heap_increase: 0
stable_memory_increase: 0
scopes: { }
scopes: {}
btreemap_get_u64_u64_v2_mem_manager:
total:
instructions: 522662446
instructions: 419339926
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down Expand Up @@ -220,10 +220,10 @@ benches:
instructions: 5070548871
heap_increase: 0
stable_memory_increase: 261
scopes: { }
scopes: {}
btreemap_insert_blob_1024_512_v2_mem_manager:
total:
instructions: 5594888576
instructions: 5399122036
heap_increase: 0
stable_memory_increase: 256
scopes: {}
Expand Down Expand Up @@ -349,13 +349,13 @@ benches:
scopes: {}
btreemap_insert_blob_8_u64:
total:
instructions: 324571604
instructions: 324578222
heap_increase: 0
stable_memory_increase: 6
scopes: {}
btreemap_insert_blob_8_u64_v2:
total:
instructions: 429525034
instructions: 429531652
heap_increase: 0
stable_memory_increase: 4
scopes: {}
Expand All @@ -379,10 +379,10 @@ benches:
scopes: {}
btreemap_insert_u64_u64_mem_manager:
total:
instructions: 830923183
instructions: 677264271
heap_increase: 0
stable_memory_increase: 0
scopes: { }
scopes: {}
btreemap_insert_u64_u64_v2:
total:
instructions: 421501977
Expand Down Expand Up @@ -559,13 +559,13 @@ benches:
scopes: {}
btreemap_remove_blob_8_u64:
total:
instructions: 425348358
instructions: 425371740
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_remove_blob_8_u64_v2:
total:
instructions: 566089549
instructions: 566112931
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand All @@ -583,13 +583,13 @@ benches:
scopes: {}
btreemap_remove_u64_u64:
total:
instructions: 491364624
instructions: 491394624
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_remove_u64_u64_v2:
total:
instructions: 609098227
instructions: 609128227
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down Expand Up @@ -631,7 +631,7 @@ benches:
scopes: {}
memory_manager_overhead:
total:
instructions: 1182056386
instructions: 1182001071
heap_increase: 0
stable_memory_increase: 8320
scopes: {}
Expand Down Expand Up @@ -661,19 +661,19 @@ benches:
scopes: {}
vec_get_blob_4_mem_manager:
total:
instructions: 12194073
instructions: 9246773
heap_increase: 0
stable_memory_increase: 0
scopes: { }
scopes: {}
vec_get_blob_64:
total:
instructions: 12960294
heap_increase: 0
stable_memory_increase: 0
scopes: { }
scopes: {}
vec_get_blob_64_mem_manager:
total:
instructions: 20656466
instructions: 17603274
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand All @@ -685,7 +685,7 @@ benches:
scopes: {}
vec_get_u64:
total:
instructions: 5390319
instructions: 5420319
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down
178 changes: 109 additions & 69 deletions src/memory_manager.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -380,14 +380,14 @@ impl<M: Memory> MemoryManagerInner<M> {
}

let mut bytes_written = 0;
for Segment { address, length } in self.bucket_iter(id, offset, src.len()) {
self.for_each_bucket(id, offset, src.len(), |Segment { address, length }| {
self.memory.write(
address.get(),
&src[bytes_written as usize..(bytes_written + length.get()) as usize],
);

bytes_written += length.get();
}
});
}

#[inline]
Expand All @@ -407,30 +407,82 @@ impl<M: Memory> MemoryManagerInner<M> {
}

let mut bytes_read: usize = 0;
for Segment { address, length } in self.bucket_iter(id, offset, count) {
self.for_each_bucket(id, offset, count, |Segment { address, length }| {
let length = length.get().try_into().expect("Length overflows usize");
self.memory
.read_unsafe(address.get(), dst.add(bytes_read), length);

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

/// Initializes a [`BucketIterator`].
fn bucket_iter(&self, MemoryId(id): MemoryId, offset: u64, length: usize) -> BucketIterator {
fn for_each_bucket(
&self,
MemoryId(id): MemoryId,
offset: u64,
length: usize,
mut func: impl FnMut(Segment),
) {
if length == 0 {
return;
}

// TODO: check length
// Get the buckets allocated to the given memory id.
let buckets = self.memory_buckets[id as usize].as_slice();

BucketIterator {
virtual_segment: Segment {
address: Address::from(offset),
length: Bytes::from(length as u64),
},
buckets,
bucket_size_in_bytes: self.bucket_size_in_bytes(),
let mut buckets = self.memory_buckets[id as usize].as_slice();
let mut virtual_segment = Segment {
address: Address::from(offset),
length: Bytes::from(length as u64),
};

let bucket_size_in_bytes = self.bucket_size_in_bytes().get();
// starting bucket index
let mut bucket_idx = (offset / bucket_size_in_bytes) as usize;
let start_offset = offset % bucket_size_in_bytes;
let remaining_bytes = bucket_size_in_bytes - start_offset;

let bucket_address =
self.bucket_address(buckets.get(bucket_idx).expect("bucket idx out of bounds"));

let first_segment_len = remaining_bytes.min(length as u64);
func(Segment {
address: bucket_address + start_offset.into(),
length: first_segment_len.into(),
});

virtual_segment.length -= first_segment_len.into();
virtual_segment.address += first_segment_len.into();

if virtual_segment.length.get() == 0 {
return;
}

for _ in 0..(virtual_segment.length.get() / self.bucket_size_in_bytes().get()) {
bucket_idx += 1;
let bucket_address =
self.bucket_address(buckets.get(bucket_idx).expect("bucket idx out of bounds"));
func(Segment {
address: bucket_address,
length: self.bucket_size_in_bytes(),
});
}

if virtual_segment.length.get() == 0 {
return;
}

bucket_idx += 1;

let bucket_address =
self.bucket_address(buckets.get(bucket_idx).expect("bucket idx out of bounds"));

func(Segment {
address: bucket_address,
length: virtual_segment.length % self.bucket_size_in_bytes(),
});
}

fn bucket_size_in_bytes(&self) -> Bytes {
Expand All @@ -447,6 +499,11 @@ impl<M: Memory> MemoryManagerInner<M> {
pub fn into_memory(self) -> M {
self.memory
}

#[inline]
fn bucket_address(&self, id: &BucketId) -> Address {
Address::from(BUCKETS_OFFSET_IN_BYTES) + self.bucket_size_in_bytes() * Bytes::from(id.0)
}
}

struct Segment {
Expand Down Expand Up @@ -477,61 +534,44 @@ struct Segment {
//
// Each of the segments above can then be translated into the real address space by looking up
// the underlying buckets' addresses in real memory.
struct BucketIterator<'a> {
virtual_segment: Segment,
buckets: &'a [BucketId],
bucket_size_in_bytes: Bytes,
}

impl Iterator for BucketIterator<'_> {
type Item = Segment;

fn next(&mut self) -> Option<Self::Item> {
if self.virtual_segment.length == Bytes::from(0u64) {
return None;
}

// Map the virtual segment's address to a real address.
let bucket_idx =
(self.virtual_segment.address.get() / self.bucket_size_in_bytes.get()) as usize;
let bucket_address = self.bucket_address(
*self
.buckets
.get(bucket_idx)
.expect("bucket idx out of bounds"),
);

let real_address = bucket_address
+ Bytes::from(self.virtual_segment.address.get() % self.bucket_size_in_bytes.get());

// Compute how many bytes are in this real segment.
let bytes_in_segment = {
let next_bucket_address = bucket_address + self.bucket_size_in_bytes;

// Write up to either the end of the bucket, or the end of the segment.
min(
Bytes::from(next_bucket_address.get() - real_address.get()),
self.virtual_segment.length,
)
};

// Update the virtual segment to exclude the portion we're about to return.
self.virtual_segment.length -= bytes_in_segment;
self.virtual_segment.address += bytes_in_segment;

Some(Segment {
address: real_address,
length: bytes_in_segment,
})
}
}

impl<'a> BucketIterator<'a> {
/// 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)
}
}
// struct BucketIterator<'a> {
// virtual_segment: Segment,
// buckets: &'a [BucketId],
// remaining_space_in_bucket: Bytes,
// bucket_size_in_bytes: Bytes,
// }
//
// impl Iterator for BucketIterator<'_> {
// type Item = Segment;
//
// fn next(&mut self) -> Option<Self::Item> {
// if self.virtual_segment.length == Bytes::from(0u64) {
// return None;
// }
//
// // Map the virtual segment's address to a real address.
// let bucket_address =
// self.bucket_address(*self.buckets.get(0).expect("bucket idx out of bounds"));
//
// let real_address = bucket_address
// + Bytes::from(self.virtual_segment.address.get() % self.bucket_size_in_bytes.get());
//
// // Compute how many bytes are in this real segment.
// // Write up to either the remaining space in current bucket, or the end of the segment.
// let bytes_in_segment = min(self.remaining_space_in_bucket, self.virtual_segment.length);
//
// // Update the virtual segment to exclude the portion we're about to return.
// self.virtual_segment.length -= bytes_in_segment;
// self.virtual_segment.address += bytes_in_segment;
// self.buckets = &self.buckets[1..];
// self.remaining_space_in_bucket = self.bucket_size_in_bytes;
//
// Some(Segment {
// address: real_address,
// length: bytes_in_segment,
// })
// }
// }

#[derive(Clone, Copy, Ord, Eq, PartialEq, PartialOrd, Debug)]
pub struct MemoryId(u8);
Expand Down

0 comments on commit 257abad

Please sign in to comment.