From efdefe1bd718f3a22e5c65ec6b65149507193704 Mon Sep 17 00:00:00 2001 From: David Frank Date: Wed, 4 Dec 2024 14:34:37 +0000 Subject: [PATCH 1/2] perf: Add bucket cache in MemoryManager Since buckets are quite big (minimum 64kb and 8MB by default), subsequent accesses will likely touch the same bucket. It's worth caching the last accessed bucket. --- canbench_results.yml | 14 ++-- src/memory_manager.rs | 181 +++++++++++++++++++++++++++++++++++++++++- src/types.rs | 8 ++ 3 files changed, 195 insertions(+), 8 deletions(-) diff --git a/canbench_results.yml b/canbench_results.yml index 1ed74729..bfa4f55d 100644 --- a/canbench_results.yml +++ b/canbench_results.yml @@ -73,7 +73,7 @@ benches: scopes: {} btreemap_get_blob_512_1024_v2_mem_manager: total: - instructions: 2624136090 + instructions: 2529776211 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -139,7 +139,7 @@ benches: scopes: {} btreemap_get_u64_u64_v2_mem_manager: total: - instructions: 421366446 + instructions: 346385206 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -223,7 +223,7 @@ benches: scopes: {} btreemap_insert_blob_1024_512_v2_mem_manager: total: - instructions: 5402984503 + instructions: 5258158583 heap_increase: 0 stable_memory_increase: 256 scopes: {} @@ -379,7 +379,7 @@ benches: scopes: {} btreemap_insert_u64_u64_mem_manager: total: - instructions: 680292499 + instructions: 565125199 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -631,7 +631,7 @@ benches: scopes: {} memory_manager_overhead: total: - instructions: 1182002161 + instructions: 1181970633 heap_increase: 0 stable_memory_increase: 8320 scopes: {} @@ -661,7 +661,7 @@ benches: scopes: {} vec_get_blob_4_mem_manager: total: - instructions: 9333723 + instructions: 7351373 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -673,7 +673,7 @@ benches: scopes: {} vec_get_blob_64_mem_manager: total: - instructions: 17664902 + instructions: 15459258 heap_increase: 0 stable_memory_increase: 0 scopes: {} diff --git a/src/memory_manager.rs b/src/memory_manager.rs index f2c0990b..d941d2e2 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -45,7 +45,7 @@ use crate::{ types::{Address, Bytes}, write, write_struct, Memory, WASM_PAGE_SIZE, }; -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; use std::rc::Rc; const MAGIC: &[u8; 3] = b"MGR"; @@ -233,6 +233,8 @@ struct MemoryManagerInner { /// A map mapping each managed memory to the bucket ids that are allocated to it. memory_buckets: Vec>, + + bucket_cache: BucketCache, } impl MemoryManagerInner { @@ -261,6 +263,7 @@ impl MemoryManagerInner { memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize], memory_buckets: vec![vec![]; MAX_NUM_MEMORIES as usize], bucket_size_in_pages, + bucket_cache: BucketCache::new(), }; mem_mgr.save_header(); @@ -302,6 +305,7 @@ impl MemoryManagerInner { bucket_size_in_pages: header.bucket_size_in_pages, memory_sizes_in_pages: header.memory_sizes_in_pages, memory_buckets, + bucket_cache: BucketCache::new(), } } @@ -374,6 +378,17 @@ impl MemoryManagerInner { } fn write(&self, id: MemoryId, offset: u64, src: &[u8]) { + if let Some(real_address) = self.bucket_cache.get( + id, + VirtualSegment { + address: offset.into(), + length: src.len().into(), + }, + ) { + self.memory.write(real_address.get(), src); + return; + } + if (offset + src.len() as u64) > self.memory_size(id) * WASM_PAGE_SIZE { panic!("{id:?}: write out of bounds"); } @@ -408,6 +423,18 @@ impl MemoryManagerInner { /// * 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) { + // First try to find the virtual segment in the cache. + if let Some(real_address) = self.bucket_cache.get( + id, + VirtualSegment { + address: offset.into(), + length: count.into(), + }, + ) { + self.memory.read_unsafe(real_address.get(), dst, count); + return; + } + if (offset + count as u64) > self.memory_size(id) * WASM_PAGE_SIZE { panic!("{id:?}: read out of bounds"); } @@ -482,8 +509,20 @@ impl MemoryManagerInner { while length > 0 { let bucket_address = self.bucket_address(buckets.get(bucket_idx).expect("bucket idx out of bounds")); + + let bucket_start = bucket_idx as u64 * bucket_size_in_bytes; let segment_len = (bucket_size_in_bytes - start_offset_in_bucket).min(length); + // Cache this bucket. + self.bucket_cache.store( + MemoryId(id), + VirtualSegment { + address: bucket_start.into(), + length: self.bucket_size_in_bytes(), + }, + bucket_address, + ); + func(RealSegment { address: bucket_address + start_offset_in_bucket.into(), length: segment_len.into(), @@ -516,11 +555,18 @@ impl MemoryManagerInner { } } +#[derive(Copy, Clone)] struct VirtualSegment { address: Address, length: Bytes, } +impl VirtualSegment { + fn contains_segment(&self, other: &VirtualSegment) -> bool { + self.address <= other.address && other.address + other.length <= self.address + self.length + } +} + struct RealSegment { address: Address, length: Bytes, @@ -547,6 +593,53 @@ fn bucket_allocations_address(id: BucketId) -> Address { Address::from(0) + Header::size() + Bytes::from(id.0) } +/// Cache which stores the last touched bucket and the corresponding real address. +/// +/// If a segment from this bucket is accessed, we can return the real address faster. +#[derive(Clone)] +struct BucketCache { + memory_id: Cell, + bucket: Cell, + /// The real address that corresponds to bucket.address + real_address: Cell
, +} + +impl BucketCache { + #[inline] + fn new() -> Self { + BucketCache { + memory_id: Cell::new(MemoryId(0)), + bucket: Cell::new(VirtualSegment { + address: Address::from(0), + length: Bytes::new(0), + }), + real_address: Cell::new(Address::from(0)), + } + } +} + +impl BucketCache { + /// Returns the real address corresponding to `virtual_segment.address` if `virtual_segment` + /// is fully contained within the cached bucket, otherwise `None`. + #[inline] + fn get(&self, memory_id: MemoryId, virtual_segment: VirtualSegment) -> Option
{ + let cached_bucket = self.bucket.get(); + let cache_hit = + self.memory_id.get() == memory_id && cached_bucket.contains_segment(&virtual_segment); + + cache_hit + .then(|| self.real_address.get() + (virtual_segment.address - cached_bucket.address)) + } + + /// Stores the mapping of a bucket to a real address. + #[inline] + fn store(&self, memory_id: MemoryId, bucket: VirtualSegment, real_address: Address) { + self.memory_id.set(memory_id); + self.bucket.set(bucket); + self.real_address.set(real_address); + } +} + #[cfg(test)] mod test { use super::*; @@ -950,4 +1043,90 @@ mod test { let expected_read = include_bytes!("memory_manager/stability_read.golden"); assert!(expected_read.as_slice() == read.as_slice()); } + + #[test] + fn bucket_cache() { + let bucket_cache = BucketCache::new(); + + // No match, nothing has been stored. + assert_eq!( + bucket_cache.get( + MemoryId::new(0), + VirtualSegment { + address: Address::from(0), + length: Bytes::from(1u64) + } + ), + None + ); + + bucket_cache.store( + MemoryId::new(22), + VirtualSegment { + address: Address::from(0), + length: Bytes::from(335u64), + }, + Address::from(983), + ); + + // Match at the beginning + assert_eq!( + bucket_cache.get( + MemoryId::new(22), + VirtualSegment { + address: Address::from(1), + length: Bytes::from(2u64) + } + ), + Some(Address::from(984)) + ); + + // Match at the end + assert_eq!( + bucket_cache.get( + MemoryId::new(22), + VirtualSegment { + address: Address::from(334), + length: Bytes::from(1u64) + } + ), + Some(Address::from(1317)) + ); + + // Match entire segment + assert_eq!( + bucket_cache.get( + MemoryId::new(22), + VirtualSegment { + address: Address::from(0), + length: Bytes::from(335u64), + } + ), + Some(Address::from(983)) + ); + + // No match (memory id is different) + assert_eq!( + bucket_cache.get( + MemoryId::new(23), + VirtualSegment { + address: Address::from(1), + length: Bytes::from(2u64) + } + ), + None + ); + + // No match - outside cached segment + assert_eq!( + bucket_cache.get( + MemoryId::new(22), + VirtualSegment { + address: Address::from(1), + length: Bytes::from(335u64) + } + ), + None + ); + } } diff --git a/src/types.rs b/src/types.rs index ffb3e3ad..0c9bbf7f 100644 --- a/src/types.rs +++ b/src/types.rs @@ -31,6 +31,14 @@ impl Add for Address { } } +impl Sub
for Address { + type Output = Bytes; + + fn sub(self, address: Address) -> Self::Output { + Bytes(self.0 - address.0) + } +} + impl Sub for Address { type Output = Self; From eb138e8b38200fadd793517e0b0a94623b0fee30 Mon Sep 17 00:00:00 2001 From: David Frank Date: Wed, 4 Dec 2024 17:21:42 +0000 Subject: [PATCH 2/2] Cache per VirtualMemory --- canbench_results.yml | 14 ++--- src/memory_manager.rs | 139 +++++++++++++++++------------------------- 2 files changed, 63 insertions(+), 90 deletions(-) diff --git a/canbench_results.yml b/canbench_results.yml index bfa4f55d..15d588c4 100644 --- a/canbench_results.yml +++ b/canbench_results.yml @@ -73,7 +73,7 @@ benches: scopes: {} btreemap_get_blob_512_1024_v2_mem_manager: total: - instructions: 2529776211 + instructions: 2517895062 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -139,7 +139,7 @@ benches: scopes: {} btreemap_get_u64_u64_v2_mem_manager: total: - instructions: 346385206 + instructions: 337472533 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -223,7 +223,7 @@ benches: scopes: {} btreemap_insert_blob_1024_512_v2_mem_manager: total: - instructions: 5258158583 + instructions: 5243334303 heap_increase: 0 stable_memory_increase: 256 scopes: {} @@ -379,7 +379,7 @@ benches: scopes: {} btreemap_insert_u64_u64_mem_manager: total: - instructions: 565125199 + instructions: 553691634 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -631,7 +631,7 @@ benches: scopes: {} memory_manager_overhead: total: - instructions: 1181970633 + instructions: 1181967369 heap_increase: 0 stable_memory_increase: 8320 scopes: {} @@ -661,7 +661,7 @@ benches: scopes: {} vec_get_blob_4_mem_manager: total: - instructions: 7351373 + instructions: 7238723 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -673,7 +673,7 @@ benches: scopes: {} vec_get_blob_64_mem_manager: total: - instructions: 15459258 + instructions: 15339702 heap_increase: 0 stable_memory_increase: 0 scopes: {} diff --git a/src/memory_manager.rs b/src/memory_manager.rs index d941d2e2..10cd9028 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -151,6 +151,7 @@ impl MemoryManager { VirtualMemory { id, memory_manager: self.inner.clone(), + cache: BucketCache::new(), } } @@ -193,6 +194,7 @@ impl Header { pub struct VirtualMemory { id: MemoryId, memory_manager: Rc>>, + cache: BucketCache, } impl Memory for VirtualMemory { @@ -205,17 +207,21 @@ impl Memory for VirtualMemory { } fn read(&self, offset: u64, dst: &mut [u8]) { - self.memory_manager.borrow().read(self.id, offset, dst) + self.memory_manager + .borrow() + .read(self.id, offset, dst, &self.cache) } unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) { self.memory_manager .borrow() - .read_unsafe(self.id, offset, dst, count) + .read_unsafe(self.id, offset, dst, count, &self.cache) } fn write(&self, offset: u64, src: &[u8]) { - self.memory_manager.borrow().write(self.id, offset, src) + self.memory_manager + .borrow() + .write(self.id, offset, src, &self.cache) } } @@ -233,8 +239,6 @@ struct MemoryManagerInner { /// A map mapping each managed memory to the bucket ids that are allocated to it. memory_buckets: Vec>, - - bucket_cache: BucketCache, } impl MemoryManagerInner { @@ -263,7 +267,6 @@ impl MemoryManagerInner { memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize], memory_buckets: vec![vec![]; MAX_NUM_MEMORIES as usize], bucket_size_in_pages, - bucket_cache: BucketCache::new(), }; mem_mgr.save_header(); @@ -305,7 +308,6 @@ impl MemoryManagerInner { bucket_size_in_pages: header.bucket_size_in_pages, memory_sizes_in_pages: header.memory_sizes_in_pages, memory_buckets, - bucket_cache: BucketCache::new(), } } @@ -377,14 +379,11 @@ impl MemoryManagerInner { old_size as i64 } - fn write(&self, id: MemoryId, offset: u64, src: &[u8]) { - if let Some(real_address) = self.bucket_cache.get( - id, - VirtualSegment { - address: offset.into(), - length: src.len().into(), - }, - ) { + fn write(&self, id: MemoryId, offset: u64, src: &[u8], bucket_cache: &BucketCache) { + if let Some(real_address) = bucket_cache.get(VirtualSegment { + address: offset.into(), + length: src.len().into(), + }) { self.memory.write(real_address.get(), src); return; } @@ -400,6 +399,7 @@ impl MemoryManagerInner { address: offset.into(), length: src.len().into(), }, + bucket_cache, |RealSegment { address, length }| { self.memory.write( address.get(), @@ -412,9 +412,9 @@ impl MemoryManagerInner { } #[inline] - fn read(&self, id: MemoryId, offset: u64, dst: &mut [u8]) { + fn read(&self, id: MemoryId, offset: u64, dst: &mut [u8], bucket_cache: &BucketCache) { // SAFETY: this is trivially safe because dst has dst.len() space. - unsafe { self.read_unsafe(id, offset, dst.as_mut_ptr(), dst.len()) } + unsafe { self.read_unsafe(id, offset, dst.as_mut_ptr(), dst.len(), bucket_cache) } } /// # Safety @@ -422,15 +422,19 @@ impl MemoryManagerInner { /// 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) { + unsafe fn read_unsafe( + &self, + id: MemoryId, + offset: u64, + dst: *mut u8, + count: usize, + bucket_cache: &BucketCache, + ) { // First try to find the virtual segment in the cache. - if let Some(real_address) = self.bucket_cache.get( - id, - VirtualSegment { - address: offset.into(), - length: count.into(), - }, - ) { + if let Some(real_address) = bucket_cache.get(VirtualSegment { + address: offset.into(), + length: count.into(), + }) { self.memory.read_unsafe(real_address.get(), dst, count); return; } @@ -446,6 +450,7 @@ impl MemoryManagerInner { address: offset.into(), length: count.into(), }, + bucket_cache, |RealSegment { address, length }| { self.memory.read_unsafe( address.get(), @@ -492,6 +497,7 @@ impl MemoryManagerInner { &self, MemoryId(id): MemoryId, virtual_segment: VirtualSegment, + bucket_cache: &BucketCache, mut func: impl FnMut(RealSegment), ) { // Get the buckets allocated to the given memory id. @@ -514,8 +520,7 @@ impl MemoryManagerInner { let segment_len = (bucket_size_in_bytes - start_offset_in_bucket).min(length); // Cache this bucket. - self.bucket_cache.store( - MemoryId(id), + bucket_cache.store( VirtualSegment { address: bucket_start.into(), length: self.bucket_size_in_bytes(), @@ -598,7 +603,6 @@ fn bucket_allocations_address(id: BucketId) -> Address { /// If a segment from this bucket is accessed, we can return the real address faster. #[derive(Clone)] struct BucketCache { - memory_id: Cell, bucket: Cell, /// The real address that corresponds to bucket.address real_address: Cell
, @@ -608,7 +612,6 @@ impl BucketCache { #[inline] fn new() -> Self { BucketCache { - memory_id: Cell::new(MemoryId(0)), bucket: Cell::new(VirtualSegment { address: Address::from(0), length: Bytes::new(0), @@ -622,19 +625,17 @@ impl BucketCache { /// Returns the real address corresponding to `virtual_segment.address` if `virtual_segment` /// is fully contained within the cached bucket, otherwise `None`. #[inline] - fn get(&self, memory_id: MemoryId, virtual_segment: VirtualSegment) -> Option
{ + fn get(&self, virtual_segment: VirtualSegment) -> Option
{ let cached_bucket = self.bucket.get(); - let cache_hit = - self.memory_id.get() == memory_id && cached_bucket.contains_segment(&virtual_segment); - cache_hit + cached_bucket + .contains_segment(&virtual_segment) .then(|| self.real_address.get() + (virtual_segment.address - cached_bucket.address)) } /// Stores the mapping of a bucket to a real address. #[inline] - fn store(&self, memory_id: MemoryId, bucket: VirtualSegment, real_address: Address) { - self.memory_id.set(memory_id); + fn store(&self, bucket: VirtualSegment, real_address: Address) { self.bucket.set(bucket); self.real_address.set(real_address); } @@ -1050,18 +1051,14 @@ mod test { // No match, nothing has been stored. assert_eq!( - bucket_cache.get( - MemoryId::new(0), - VirtualSegment { - address: Address::from(0), - length: Bytes::from(1u64) - } - ), + bucket_cache.get(VirtualSegment { + address: Address::from(0), + length: Bytes::from(1u64) + }), None ); bucket_cache.store( - MemoryId::new(22), VirtualSegment { address: Address::from(0), length: Bytes::from(335u64), @@ -1071,61 +1068,37 @@ mod test { // Match at the beginning assert_eq!( - bucket_cache.get( - MemoryId::new(22), - VirtualSegment { - address: Address::from(1), - length: Bytes::from(2u64) - } - ), + bucket_cache.get(VirtualSegment { + address: Address::from(1), + length: Bytes::from(2u64) + }), Some(Address::from(984)) ); // Match at the end assert_eq!( - bucket_cache.get( - MemoryId::new(22), - VirtualSegment { - address: Address::from(334), - length: Bytes::from(1u64) - } - ), + bucket_cache.get(VirtualSegment { + address: Address::from(334), + length: Bytes::from(1u64) + }), Some(Address::from(1317)) ); // Match entire segment assert_eq!( - bucket_cache.get( - MemoryId::new(22), - VirtualSegment { - address: Address::from(0), - length: Bytes::from(335u64), - } - ), + bucket_cache.get(VirtualSegment { + address: Address::from(0), + length: Bytes::from(335u64), + }), Some(Address::from(983)) ); - // No match (memory id is different) - assert_eq!( - bucket_cache.get( - MemoryId::new(23), - VirtualSegment { - address: Address::from(1), - length: Bytes::from(2u64) - } - ), - None - ); - // No match - outside cached segment assert_eq!( - bucket_cache.get( - MemoryId::new(22), - VirtualSegment { - address: Address::from(1), - length: Bytes::from(335u64) - } - ), + bucket_cache.get(VirtualSegment { + address: Address::from(1), + length: Bytes::from(335u64) + }), None ); }