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.

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

fix(allocator): Fix allocator & add benchmark #9234

Merged
merged 22 commits into from
Jul 14, 2024
3 changes: 3 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 12 additions & 0 deletions crates/swc_allocator/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,3 +24,15 @@ rkyv = { workspace = true, optional = true }
scoped-tls = { workspace = true }
serde = { workspace = true, optional = true }
serde_derive = { workspace = true, optional = true }


[dev-dependencies]
criterion = { workspace = true }

codspeed-criterion-compat = { workspace = true }
swc_malloc = { version = "0.5.10", path = "../swc_malloc" }


[[bench]]
harness = false
name = "bench"
92 changes: 92 additions & 0 deletions crates/swc_allocator/benches/bench.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
extern crate swc_malloc;

use codspeed_criterion_compat::{black_box, criterion_group, criterion_main, Bencher, Criterion};
use swc_allocator::{FastAlloc, MemorySpace};

fn bench_alloc(c: &mut Criterion) {
fn direct_alloc_std(b: &mut Bencher, times: usize) {
b.iter(|| {
let mut buf = std::vec::Vec::new();
for i in 0..times {
let item: std::boxed::Box<usize> = black_box(std::boxed::Box::new(black_box(i)));
buf.push(item);
}
})
}

fn direct_alloc_no_scope(b: &mut Bencher, times: usize) {
b.iter(|| {
let mut vec = swc_allocator::vec::Vec::new();
for i in 0..times {
let item: swc_allocator::boxed::Box<usize> =
black_box(swc_allocator::boxed::Box::new(black_box(i)));
vec.push(item);
}
})
}

fn fast_alloc_no_scope(b: &mut Bencher, times: usize) {
b.iter(|| {
let allocator = FastAlloc::default();

let mut vec = allocator.vec();
for i in 0..times {
let item: swc_allocator::boxed::Box<usize> =
black_box(allocator.alloc(black_box(i)));
vec.push(item);
}
})
}

fn direct_alloc_scoped(b: &mut Bencher, times: usize) {
b.iter(|| {
let allocator = MemorySpace::default();

allocator.scope(|| {
let mut vec = swc_allocator::vec::Vec::new();

for i in 0..times {
let item: swc_allocator::boxed::Box<usize> =
black_box(swc_allocator::boxed::Box::new(black_box(i)));
vec.push(item);
}
});
})
}

fn fast_alloc_scoped(b: &mut Bencher, times: usize) {
b.iter(|| {
MemorySpace::default().scope(|| {
let allocator = FastAlloc::default();

let mut vec = allocator.vec();

for i in 0..times {
let item: swc_allocator::boxed::Box<usize> =
black_box(allocator.alloc(black_box(i)));
vec.push(item);
}
});
})
}

c.bench_function("common/allocator/alloc/std/100000", |b| {
direct_alloc_std(b, 100000)
});
c.bench_function("common/allocator/alloc/no-scope/100000", |b| {
direct_alloc_no_scope(b, 100000)
});
c.bench_function("common/allocator/alloc/scoped/100000", |b| {
direct_alloc_scoped(b, 100000)
});

c.bench_function("common/allocator/alloc/cached-no-scope/100000", |b| {
fast_alloc_no_scope(b, 100000)
});
c.bench_function("common/allocator/alloc/cached-scoped/100000", |b| {
fast_alloc_scoped(b, 100000)
});
}

criterion_group!(benches, bench_alloc);
criterion_main!(benches);
49 changes: 27 additions & 22 deletions crates/swc_allocator/src/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,30 @@ use std::{alloc::Layout, ptr::NonNull};
use allocator_api2::alloc::Global;
use scoped_tls::scoped_thread_local;

use crate::Allocator;
use crate::{FastAlloc, MemorySpace};

scoped_thread_local!(pub(crate) static ALLOC: Allocator);
scoped_thread_local!(pub(crate) static ALLOC: MemorySpace);

#[derive(Debug, Clone, Copy, Default)]
pub struct SwcAlloc;
#[derive(Debug, Clone, Copy)]
pub struct SwcAlloc {
pub(crate) is_arena_mode: bool,
}

impl Default for FastAlloc {
fn default() -> Self {
Self {
is_arena_mode: ALLOC.is_set(),
}
}
}

impl Default for SwcAlloc {
fn default() -> Self {
SwcAlloc {
is_arena_mode: ALLOC.is_set(),
}
}
}

impl SwcAlloc {
/// `true` is passed to `f` if the box is allocated with a custom allocator.
Expand All @@ -27,21 +45,8 @@ impl SwcAlloc {
}
}

/// Set the last bit to 1
fn mark_ptr_as_arena_mode(ptr: NonNull<[u8]>) -> NonNull<[u8]> {
let (mut raw_ptr, metadata) = ptr_meta::PtrExt::to_raw_parts(ptr.as_ptr());

raw_ptr = (raw_ptr as usize | 1) as *mut ();

unsafe {
// Safety:
NonNull::new_unchecked(ptr_meta::from_raw_parts_mut(raw_ptr, metadata))
}
}

fn is_ptr_in_arena_mode(ptr: NonNull<u8>) -> bool {
let ptr = ptr.as_ptr() as usize;
ptr & 1 == 1
ptr
}

unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
Expand Down Expand Up @@ -73,7 +78,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
}

unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
if is_ptr_in_arena_mode(ptr) {
if self.is_arena_mode {
debug_assert!(
ALLOC.is_set(),
"Deallocating a pointer allocated with arena mode with a non-arena mode allocator"
Expand All @@ -96,7 +101,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
old_layout: Layout,
new_layout: Layout,
) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
if is_ptr_in_arena_mode(ptr) {
if self.is_arena_mode {
debug_assert!(
ALLOC.is_set(),
"Growing a pointer allocated with arena mode with a non-arena mode allocator"
Expand All @@ -114,7 +119,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
old_layout: Layout,
new_layout: Layout,
) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
if is_ptr_in_arena_mode(ptr) {
if self.is_arena_mode {
debug_assert!(
ALLOC.is_set(),
"Growing a pointer allocated with arena mode with a non-arena mode allocator"
Expand All @@ -132,7 +137,7 @@ unsafe impl allocator_api2::alloc::Allocator for SwcAlloc {
old_layout: Layout,
new_layout: Layout,
) -> Result<NonNull<[u8]>, allocator_api2::alloc::AllocError> {
if is_ptr_in_arena_mode(ptr) {
if self.is_arena_mode {
debug_assert!(
ALLOC.is_set(),
"Shrinking a pointer allocated with arena mode with a non-arena mode allocator"
Expand Down
18 changes: 15 additions & 3 deletions crates/swc_allocator/src/boxed/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use std::{
pin::Pin,
};

use crate::alloc::SwcAlloc;
use crate::{alloc::SwcAlloc, FastAlloc};

#[cfg(feature = "rkyv")]
mod rkyv;
Expand Down Expand Up @@ -54,7 +54,10 @@ impl<T> Box<T> {
/// See [`std::boxed::Box::new`].
#[inline(always)]
pub fn new(value: T) -> Self {
Self(allocator_api2::boxed::Box::new_in(value, SwcAlloc))
Self(allocator_api2::boxed::Box::new_in(
value,
SwcAlloc::default(),
))
}

/// Moves the value out of the box.
Expand Down Expand Up @@ -106,7 +109,10 @@ impl<T: ?Sized> Box<T> {
/// [memory layout]: self#memory-layout
/// [`Layout`]: crate::Layout
pub unsafe fn from_raw(raw: *mut T) -> Self {
Self(allocator_api2::boxed::Box::from_raw_in(raw, SwcAlloc))
Self(allocator_api2::boxed::Box::from_raw_in(
raw,
SwcAlloc::default(),
))
}

/// Consumes the `Box`, returning a wrapped raw pointer.
Expand Down Expand Up @@ -621,3 +627,9 @@ where
self.0.len()
}
}

impl FastAlloc {
pub fn alloc<T>(self, t: T) -> Box<T> {
Box(allocator_api2::boxed::Box::new_in(t, self.swc_alloc()))
}
}
24 changes: 19 additions & 5 deletions crates/swc_allocator/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

#![allow(clippy::needless_doctest_main)]

use alloc::SwcAlloc;
use std::ops::{Deref, DerefMut};

use bumpalo::Bump;
Expand All @@ -14,12 +15,25 @@ mod alloc;
pub mod boxed;
pub mod vec;

#[derive(Debug, Clone, Copy)]
pub struct FastAlloc {
is_arena_mode: bool,
}

impl FastAlloc {
fn swc_alloc(self) -> SwcAlloc {
SwcAlloc {
is_arena_mode: self.is_arena_mode,
}
}
}

#[derive(Default)]
pub struct Allocator {
pub struct MemorySpace {
alloc: Bump,
}

impl Allocator {
impl MemorySpace {
/// Invokes `f` in a scope where the allocations are done in this allocator.
#[inline(always)]
pub fn scope<F, R>(&self, f: F) -> R
Expand All @@ -30,21 +44,21 @@ impl Allocator {
}
}

impl From<Bump> for Allocator {
impl From<Bump> for MemorySpace {
fn from(alloc: Bump) -> Self {
Self { alloc }
}
}

impl Deref for Allocator {
impl Deref for MemorySpace {
type Target = Bump;

fn deref(&self) -> &Bump {
&self.alloc
}
}

impl DerefMut for Allocator {
impl DerefMut for MemorySpace {
fn deref_mut(&mut self) -> &mut Bump {
&mut self.alloc
}
Expand Down
25 changes: 21 additions & 4 deletions crates/swc_allocator/src/vec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use std::ops::{Deref, DerefMut};
#[cfg(feature = "rkyv")]
mod rkyv;

use crate::{alloc::SwcAlloc, boxed::Box};
use crate::{alloc::SwcAlloc, boxed::Box, FastAlloc};

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
Expand All @@ -20,7 +20,8 @@ impl<T> Vec<T> {

pub fn with_capacity(capacity: usize) -> Self {
Self(allocator_api2::vec::Vec::with_capacity_in(
capacity, SwcAlloc,
capacity,
SwcAlloc::default(),
))
}

Expand Down Expand Up @@ -163,7 +164,10 @@ impl<T> Vec<T> {
/// ```
pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
Self(allocator_api2::vec::Vec::from_raw_parts_in(
ptr, length, capacity, SwcAlloc,
ptr,
length,
capacity,
SwcAlloc::default(),
))
}
}
Expand All @@ -184,7 +188,7 @@ impl<T> DerefMut for Vec<T> {

impl<T> Default for Vec<T> {
fn default() -> Self {
Self(allocator_api2::vec::Vec::new_in(SwcAlloc))
Self(allocator_api2::vec::Vec::new_in(SwcAlloc::default()))
}
}

Expand Down Expand Up @@ -239,3 +243,16 @@ impl<T> Extend<T> for Vec<T> {
self.0.extend(iter)
}
}

impl FastAlloc {
pub fn vec<T>(self) -> Vec<T> {
Vec(allocator_api2::vec::Vec::new_in(self.swc_alloc()))
}

pub fn vec_with_capacity<T>(self, capacity: usize) -> Vec<T> {
Vec(allocator_api2::vec::Vec::with_capacity_in(
capacity,
self.swc_alloc(),
))
}
}
Loading
Loading