From 67b1d2aec171b4971c8133c459edde277222d66c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Fri, 3 Mar 2023 06:20:24 +0100 Subject: [PATCH 1/3] Optimize DroplessArena arena allocation --- compiler/rustc_arena/src/lib.rs | 70 ++++++++++++++++++++++++++------- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index f4900ece1cc..7ab195a7c2d 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -11,6 +11,7 @@ html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/", test(no_crate_inject, attr(deny(warnings))) )] +#![feature(core_intrinsics)] #![feature(dropck_eyepatch)] #![feature(new_uninit)] #![feature(maybe_uninit_slice)] @@ -30,11 +31,11 @@ use smallvec::SmallVec; use std::alloc::Layout; use std::cell::{Cell, RefCell}; -use std::cmp; use std::marker::PhantomData; use std::mem::{self, MaybeUninit}; use std::ptr::{self, NonNull}; use std::slice; +use std::{cmp, intrinsics}; #[inline(never)] #[cold] @@ -363,6 +364,20 @@ unsafe impl<#[may_dangle] T> Drop for TypedArena { unsafe impl Send for TypedArena {} +#[inline(always)] +fn align_down(val: usize, align: usize) -> usize { + assert!(align.is_power_of_two()); + val & !(align - 1) +} + +#[inline(always)] +fn align(val: usize, align: usize) -> usize { + assert!(align.is_power_of_two()); + (val + align - 1) & !(align - 1) +} + +const DROPLESS_ALIGNMENT: usize = mem::align_of::(); + /// An arena that can hold objects of multiple different types that impl `Copy` /// and/or satisfy `!mem::needs_drop`. pub struct DroplessArena { @@ -395,8 +410,6 @@ impl Default for DroplessArena { } impl DroplessArena { - #[inline(never)] - #[cold] fn grow(&self, additional: usize) { unsafe { let mut chunks = self.chunks.borrow_mut(); @@ -418,11 +431,30 @@ impl DroplessArena { let mut chunk = ArenaChunk::new(new_cap); self.start.set(chunk.start()); - self.end.set(chunk.end()); + + // Align the end to DROPLESS_ALIGNMENT + let end = align_down(chunk.end().addr(), DROPLESS_ALIGNMENT); + // Make sure we don't go past `start` + let end = cmp::max(chunk.start().addr(), end); + self.end.set(chunk.end().with_addr(end)); + chunks.push(chunk); } } + #[inline(never)] + #[cold] + fn grow_and_alloc_raw(&self, layout: Layout) -> *mut u8 { + self.grow(layout.size()); + self.alloc_raw(layout) + } + + #[inline(never)] + #[cold] + fn grow_and_alloc(&self) -> *mut u8 { + self.grow_and_alloc_raw(Layout::new::()) + } + /// Allocates a byte slice with specified layout from the current memory /// chunk. Returns `None` if there is no free space left to satisfy the /// request. @@ -432,10 +464,13 @@ impl DroplessArena { let old_end = self.end.get(); let end = old_end.addr(); - let align = layout.align(); - let bytes = layout.size(); + // Align allocated bytes so that `self.end` stays aligned to DROPLESS_ALIGNMENT + let bytes = align(layout.size(), DROPLESS_ALIGNMENT); - let new_end = end.checked_sub(bytes)? & !(align - 1); + // Tell LLVM that `end` is aligned to DROPLESS_ALIGNMENT + unsafe { intrinsics::assume(end == align_down(end, DROPLESS_ALIGNMENT)) }; + + let new_end = align_down(end.checked_sub(bytes)?, layout.align()); if start <= new_end { let new_end = old_end.with_addr(new_end); self.end.set(new_end); @@ -448,21 +483,26 @@ impl DroplessArena { #[inline] pub fn alloc_raw(&self, layout: Layout) -> *mut u8 { assert!(layout.size() != 0); - loop { - if let Some(a) = self.alloc_raw_without_grow(layout) { - break a; - } - // No free space left. Allocate a new chunk to satisfy the request. - // On failure the grow will panic or abort. - self.grow(layout.size()); + if let Some(a) = self.alloc_raw_without_grow(layout) { + return a; } + // No free space left. Allocate a new chunk to satisfy the request. + // On failure the grow will panic or abort. + self.grow_and_alloc_raw(layout) } #[inline] pub fn alloc(&self, object: T) -> &mut T { assert!(!mem::needs_drop::()); + assert!(mem::size_of::() != 0); - let mem = self.alloc_raw(Layout::for_value::(&object)) as *mut T; + let mem = if let Some(a) = self.alloc_raw_without_grow(Layout::for_value::(&object)) { + a + } else { + // No free space left. Allocate a new chunk to satisfy the request. + // On failure the grow will panic or abort. + self.grow_and_alloc::() + } as *mut T; unsafe { // Write into uninitialized memory. From d5d72168fa325e3cbd893803c46564974e67a517 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Sun, 5 Mar 2023 06:26:51 +0100 Subject: [PATCH 2/3] Allocate extra space to account for alignment losses --- compiler/rustc_arena/src/lib.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index 7ab195a7c2d..108476a499f 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -410,7 +410,11 @@ impl Default for DroplessArena { } impl DroplessArena { - fn grow(&self, additional: usize) { + fn grow(&self, layout: Layout) { + // Add some padding so we can align `self.end` while + // stilling fitting in a `layout` allocation. + let additional = layout.size() + cmp::max(DROPLESS_ALIGNMENT, layout.align()) - 1; + unsafe { let mut chunks = self.chunks.borrow_mut(); let mut new_cap; @@ -429,7 +433,7 @@ impl DroplessArena { // Also ensure that this chunk can fit `additional`. new_cap = cmp::max(additional, new_cap); - let mut chunk = ArenaChunk::new(new_cap); + let mut chunk = ArenaChunk::new(align(new_cap, PAGE)); self.start.set(chunk.start()); // Align the end to DROPLESS_ALIGNMENT @@ -445,8 +449,8 @@ impl DroplessArena { #[inline(never)] #[cold] fn grow_and_alloc_raw(&self, layout: Layout) -> *mut u8 { - self.grow(layout.size()); - self.alloc_raw(layout) + self.grow(layout); + self.alloc_raw_without_grow(layout).unwrap() } #[inline(never)] From 6f86591b8795fc45aad3dd59863c3b3afd246506 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Mon, 14 Aug 2023 21:29:19 +0200 Subject: [PATCH 3/3] Address comments --- compiler/rustc_arena/src/lib.rs | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index 108476a499f..e45b7c154fa 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -366,16 +366,18 @@ unsafe impl Send for TypedArena {} #[inline(always)] fn align_down(val: usize, align: usize) -> usize { - assert!(align.is_power_of_two()); + debug_assert!(align.is_power_of_two()); val & !(align - 1) } #[inline(always)] -fn align(val: usize, align: usize) -> usize { - assert!(align.is_power_of_two()); +fn align_up(val: usize, align: usize) -> usize { + debug_assert!(align.is_power_of_two()); (val + align - 1) & !(align - 1) } +// Pointer alignment is common in compiler types, so keep `DroplessArena` aligned to them +// to optimize away alignment code. const DROPLESS_ALIGNMENT: usize = mem::align_of::(); /// An arena that can hold objects of multiple different types that impl `Copy` @@ -390,6 +392,8 @@ pub struct DroplessArena { /// start. (This is slightly simpler and faster than allocating upwards, /// see .) /// When this pointer crosses the start pointer, a new chunk is allocated. + /// + /// This is kept aligned to DROPLESS_ALIGNMENT. end: Cell<*mut u8>, /// A vector of arena chunks. @@ -433,13 +437,16 @@ impl DroplessArena { // Also ensure that this chunk can fit `additional`. new_cap = cmp::max(additional, new_cap); - let mut chunk = ArenaChunk::new(align(new_cap, PAGE)); + let mut chunk = ArenaChunk::new(align_up(new_cap, PAGE)); self.start.set(chunk.start()); // Align the end to DROPLESS_ALIGNMENT let end = align_down(chunk.end().addr(), DROPLESS_ALIGNMENT); - // Make sure we don't go past `start` - let end = cmp::max(chunk.start().addr(), end); + + // Make sure we don't go past `start`. This should not happen since the allocation + // should be at least DROPLESS_ALIGNMENT - 1 bytes. + debug_assert!(chunk.start().addr() <= end); + self.end.set(chunk.end().with_addr(end)); chunks.push(chunk); @@ -469,7 +476,7 @@ impl DroplessArena { let end = old_end.addr(); // Align allocated bytes so that `self.end` stays aligned to DROPLESS_ALIGNMENT - let bytes = align(layout.size(), DROPLESS_ALIGNMENT); + let bytes = align_up(layout.size(), DROPLESS_ALIGNMENT); // Tell LLVM that `end` is aligned to DROPLESS_ALIGNMENT unsafe { intrinsics::assume(end == align_down(end, DROPLESS_ALIGNMENT)) }; @@ -477,6 +484,8 @@ impl DroplessArena { let new_end = align_down(end.checked_sub(bytes)?, layout.align()); if start <= new_end { let new_end = old_end.with_addr(new_end); + // `new_end` is aligned to DROPLESS_ALIGNMENT as `align_down` preserves alignment + // as both `end` and `bytes` are already aligned to DROPLESS_ALIGNMENT. self.end.set(new_end); Some(new_end) } else {