Skip to content

Commit bfd6b9a

Browse files
Golovanov399Golovanov399
authored
[perf] Make set_range and range_vec faster via unsafe (#1340)
* Make `set_range` and `range_vec` less safe but hopefully faster * Polish a little * Reduce code duplication
1 parent 271038b commit bfd6b9a

File tree

1 file changed

+103
-115
lines changed

1 file changed

+103
-115
lines changed

‎crates/vm/src/system/memory/paged_vec.rs‎

Lines changed: 103 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,81 @@ pub(crate) struct PagedVec<T, const PAGE_SIZE: usize> {
1313
pages: Vec<Option<Vec<T>>>,
1414
}
1515

16+
// ------------------------------------------------------------------
17+
// Common Helper Functions
18+
// These functions encapsulate the common logic for copying ranges
19+
// across pages, both for read-only and read-write (set) cases.
20+
impl<T: Default + Clone, const PAGE_SIZE: usize> PagedVec<T, PAGE_SIZE> {
21+
// Copies a range of length `len` starting at index `start`
22+
// into the memory pointed to by `dst`. If the relevant page is not
23+
// initialized, fills that portion with T::default().
24+
fn read_range_generic(&self, start: usize, len: usize, dst: *mut T) {
25+
let start_page = start / PAGE_SIZE;
26+
let end_page = (start + len - 1) / PAGE_SIZE;
27+
unsafe {
28+
if start_page == end_page {
29+
let offset = start % PAGE_SIZE;
30+
if let Some(page) = self.pages[start_page].as_ref() {
31+
ptr::copy_nonoverlapping(page.as_ptr().add(offset), dst, len);
32+
} else {
33+
std::slice::from_raw_parts_mut(dst, len).fill(T::default());
34+
}
35+
} else {
36+
let offset = start % PAGE_SIZE;
37+
let first_part = PAGE_SIZE - offset;
38+
if let Some(page) = self.pages[start_page].as_ref() {
39+
ptr::copy_nonoverlapping(page.as_ptr().add(offset), dst, first_part);
40+
} else {
41+
std::slice::from_raw_parts_mut(dst, first_part).fill(T::default());
42+
}
43+
let second_part = len - first_part;
44+
if let Some(page) = self.pages[end_page].as_ref() {
45+
ptr::copy_nonoverlapping(page.as_ptr(), dst.add(first_part), second_part);
46+
} else {
47+
std::slice::from_raw_parts_mut(dst.add(first_part), second_part)
48+
.fill(T::default());
49+
}
50+
}
51+
}
52+
}
53+
54+
// Updates a range of length `len` starting at index `start` with new values.
55+
// It copies the current values into the memory pointed to by `dst`
56+
// and then writes the new values into the underlying pages,
57+
// allocating pages (with defaults) if necessary.
58+
fn set_range_generic(&mut self, start: usize, len: usize, new: *const T, dst: *mut T) {
59+
let start_page = start / PAGE_SIZE;
60+
let end_page = (start + len - 1) / PAGE_SIZE;
61+
unsafe {
62+
if start_page == end_page {
63+
let offset = start % PAGE_SIZE;
64+
let page =
65+
self.pages[start_page].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]);
66+
ptr::copy_nonoverlapping(page.as_ptr().add(offset), dst, len);
67+
ptr::copy_nonoverlapping(new, page.as_mut_ptr().add(offset), len);
68+
} else {
69+
let offset = start % PAGE_SIZE;
70+
let first_part = PAGE_SIZE - offset;
71+
{
72+
let page =
73+
self.pages[start_page].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]);
74+
ptr::copy_nonoverlapping(page.as_ptr().add(offset), dst, first_part);
75+
ptr::copy_nonoverlapping(new, page.as_mut_ptr().add(offset), first_part);
76+
}
77+
let second_part = len - first_part;
78+
{
79+
let page =
80+
self.pages[end_page].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]);
81+
ptr::copy_nonoverlapping(page.as_ptr(), dst.add(first_part), second_part);
82+
ptr::copy_nonoverlapping(new.add(first_part), page.as_mut_ptr(), second_part);
83+
}
84+
}
85+
}
86+
}
87+
}
88+
89+
// ------------------------------------------------------------------
90+
// Implementation for types requiring Default + Clone
1691
impl<T: Default + Clone, const PAGE_SIZE: usize> PagedVec<T, PAGE_SIZE> {
1792
pub fn new(num_pages: usize) -> Self {
1893
Self {
@@ -47,33 +122,32 @@ impl<T: Default + Clone, const PAGE_SIZE: usize> PagedVec<T, PAGE_SIZE> {
47122

48123
#[inline(always)]
49124
pub fn range_vec(&self, range: Range<usize>) -> Vec<T> {
50-
let mut result = Vec::with_capacity(range.len());
51-
for page_idx in (range.start / PAGE_SIZE)..range.end.div_ceil(PAGE_SIZE) {
52-
let in_page_start = range.start.saturating_sub(page_idx * PAGE_SIZE);
53-
let in_page_end = (range.end - page_idx * PAGE_SIZE).min(PAGE_SIZE);
54-
if let Some(page) = self.pages[page_idx].as_ref() {
55-
result.extend(page[in_page_start..in_page_end].iter().cloned());
56-
} else {
57-
result.extend(vec![T::default(); in_page_end - in_page_start]);
58-
}
125+
let len = range.end - range.start;
126+
// Create a vector for uninitialized values.
127+
let mut result: Vec<MaybeUninit<T>> = Vec::with_capacity(len);
128+
// SAFETY: We set the length and then initialize every element via read_range_generic.
129+
unsafe {
130+
result.set_len(len);
131+
self.read_range_generic(range.start, len, result.as_mut_ptr() as *mut T);
132+
std::mem::transmute::<Vec<MaybeUninit<T>>, Vec<T>>(result)
59133
}
60-
result
61134
}
62135

63136
pub fn set_range(&mut self, range: Range<usize>, values: &[T]) -> Vec<T> {
64-
let mut result = Vec::with_capacity(range.len());
65-
let mut values = values.iter();
66-
for page_idx in (range.start / PAGE_SIZE)..range.end.div_ceil(PAGE_SIZE) {
67-
let in_page_start = range.start.saturating_sub(page_idx * PAGE_SIZE);
68-
let in_page_end = (range.end - page_idx * PAGE_SIZE).min(PAGE_SIZE);
69-
let page = self.pages[page_idx].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]);
70-
result.extend(
71-
page[in_page_start..in_page_end]
72-
.iter_mut()
73-
.map(|x| std::mem::replace(x, values.next().unwrap().clone())),
137+
let len = range.end - range.start;
138+
assert_eq!(values.len(), len);
139+
let mut result: Vec<MaybeUninit<T>> = Vec::with_capacity(len);
140+
// SAFETY: We will write to every element in result via set_range_generic.
141+
unsafe {
142+
result.set_len(len);
143+
self.set_range_generic(
144+
range.start,
145+
len,
146+
values.as_ptr(),
147+
result.as_mut_ptr() as *mut T,
74148
);
149+
std::mem::transmute::<Vec<MaybeUninit<T>>, Vec<T>>(result)
75150
}
76-
result
77151
}
78152

79153
pub fn memory_size(&self) -> usize {
@@ -85,112 +159,26 @@ impl<T: Default + Clone, const PAGE_SIZE: usize> PagedVec<T, PAGE_SIZE> {
85159
}
86160
}
87161

162+
// ------------------------------------------------------------------
163+
// Implementation for types requiring Default + Copy
88164
impl<T: Default + Copy, const PAGE_SIZE: usize> PagedVec<T, PAGE_SIZE> {
89165
#[inline(always)]
90166
pub fn range_array<const N: usize>(&self, from: usize) -> [T; N] {
91-
// Step 1: Create an uninitialized array of MaybeUninit<T>
167+
// Create an uninitialized array of MaybeUninit<T>
92168
let mut result: [MaybeUninit<T>; N] = unsafe {
93169
// SAFETY: An uninitialized `[MaybeUninit<T>; N]` is valid.
94170
MaybeUninit::uninit().assume_init()
95171
};
96-
97-
// Step 2: Get a mutable slice of T references from the MaybeUninit array.
98-
let result_slice = unsafe {
99-
// SAFETY: We are converting a pointer from MaybeUninit<T> to T.
100-
// This is safe because we will fully initialize every element before reading.
101-
std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut T, N)
102-
};
103-
104-
let start_page = from / PAGE_SIZE;
105-
let end_page = (from + N - 1) / PAGE_SIZE;
106-
107-
if start_page == end_page {
108-
if let Some(page) = self.pages[start_page].as_ref() {
109-
// Copy data from the page into result_slice.
110-
let page_offset = from - start_page * PAGE_SIZE;
111-
let src = &page[page_offset..page_offset + N];
112-
result_slice.copy_from_slice(src);
113-
} else {
114-
// If no page available, fill with default values.
115-
result_slice.fill(T::default());
116-
}
117-
} else {
118-
debug_assert!(start_page + 1 == end_page);
119-
let first_part = PAGE_SIZE - (from - start_page * PAGE_SIZE);
120-
if let Some(page) = self.pages[start_page].as_ref() {
121-
let page_offset = from - start_page * PAGE_SIZE;
122-
let src = &page[page_offset..];
123-
result_slice[..first_part].copy_from_slice(src);
124-
} else {
125-
result_slice[..first_part].fill(T::default());
126-
}
127-
if let Some(page) = self.pages[end_page].as_ref() {
128-
let second_part = N - first_part;
129-
let src = &page[0..second_part];
130-
result_slice[first_part..].copy_from_slice(src);
131-
} else {
132-
result_slice[first_part..].fill(T::default());
133-
}
134-
}
135-
136-
// Step 4: Convert the fully initialized array of MaybeUninit<T> to [T; N].
137-
// SAFETY: We have initialized every element of `result`.
138-
unsafe {
139-
// Transmute the array; at this point each element is initialized.
140-
ptr::read(&result as *const _ as *const [T; N])
141-
}
172+
self.read_range_generic(from, N, result.as_mut_ptr() as *mut T);
173+
// SAFETY: All elements have been initialized.
174+
unsafe { ptr::read(&result as *const _ as *const [T; N]) }
142175
}
143176

144177
#[inline(always)]
145178
pub fn set_range_array<const N: usize>(&mut self, from: usize, values: &[T; N]) -> [T; N] {
146-
// Step 1: Create an uninitialized array for old values.
179+
// Create an uninitialized array for old values.
147180
let mut result: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };
148-
let result_slice = unsafe {
149-
// SAFETY: We will fully initialize `result_slice` before reading.
150-
std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut T, N)
151-
};
152-
153-
let start_page = from / PAGE_SIZE;
154-
let end_page = (from + N - 1) / PAGE_SIZE;
155-
156-
if start_page == end_page {
157-
// Ensure the page exists; if not, allocate and fill with defaults.
158-
if self.pages[start_page].is_none() {
159-
self.pages[start_page] = Some(vec![T::default(); PAGE_SIZE]);
160-
}
161-
let page = self.pages[start_page].as_mut().unwrap();
162-
let page_offset = from - start_page * PAGE_SIZE;
163-
164-
// Copy old values from the page into result_slice.
165-
result_slice.copy_from_slice(&page[page_offset..page_offset + N]);
166-
// Write the new values from input into the page.
167-
page[page_offset..page_offset + N].copy_from_slice(&values[..]);
168-
} else {
169-
debug_assert!(start_page + 1 == end_page);
170-
let first_part = PAGE_SIZE - (from - start_page * PAGE_SIZE);
171-
172-
// Handle the first page.
173-
if self.pages[start_page].is_none() {
174-
self.pages[start_page] = Some(vec![T::default(); PAGE_SIZE]);
175-
}
176-
let page0 = self.pages[start_page].as_mut().unwrap();
177-
let page_offset = from - start_page * PAGE_SIZE;
178-
179-
result_slice[..first_part].copy_from_slice(&page0[page_offset..]);
180-
page0[page_offset..].copy_from_slice(&values[..first_part]);
181-
182-
// Handle the second page.
183-
let second_part = N - first_part;
184-
if self.pages[end_page].is_none() {
185-
self.pages[end_page] = Some(vec![T::default(); PAGE_SIZE]);
186-
}
187-
let page1 = self.pages[end_page].as_mut().unwrap();
188-
189-
result_slice[first_part..].copy_from_slice(&page1[0..second_part]);
190-
page1[0..second_part].copy_from_slice(&values[first_part..]);
191-
}
192-
193-
// Step 4: Convert the fully initialized result array to [T; N].
181+
self.set_range_generic(from, N, values.as_ptr(), result.as_mut_ptr() as *mut T);
194182
unsafe { ptr::read(&result as *const _ as *const [T; N]) }
195183
}
196184
}

0 commit comments

Comments
 (0)