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

miri: optimize zeroed alloc #136035

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions compiler/rustc_const_eval/src/const_eval/machine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,7 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeMachine<'tcx> {
Size::from_bytes(size),
align,
interpret::MemoryKind::Machine(MemoryKind::Heap),
false,
)?;
ecx.write_pointer(ptr, dest)?;
}
Expand Down
8 changes: 5 additions & 3 deletions compiler/rustc_const_eval/src/interpret/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -230,11 +230,12 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
size: Size,
align: Align,
kind: MemoryKind<M::MemoryKind>,
zero_init: bool,
) -> InterpResult<'tcx, Pointer<M::Provenance>> {
let alloc = if M::PANIC_ON_ALLOC_FAIL {
Allocation::uninit(size, align)
Allocation::new(size, align, zero_init)
} else {
Allocation::try_uninit(size, align)?
Allocation::try_new(size, align, zero_init)?
};
self.insert_allocation(alloc, kind)
}
Expand Down Expand Up @@ -277,6 +278,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
new_size: Size,
new_align: Align,
kind: MemoryKind<M::MemoryKind>,
zero_init: bool,
) -> InterpResult<'tcx, Pointer<M::Provenance>> {
let (alloc_id, offset, _prov) = self.ptr_get_alloc_id(ptr, 0)?;
if offset.bytes() != 0 {
Expand All @@ -289,7 +291,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {

// For simplicities' sake, we implement reallocate as "alloc, copy, dealloc".
// This happens so rarely, the perf advantage is outweighed by the maintenance cost.
let new_ptr = self.allocate_ptr(new_size, new_align, kind)?;
let new_ptr = self.allocate_ptr(new_size, new_align, kind, zero_init)?;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since we'll memcopy right over the alloc, zeroing is not needed for realloc, just use false here

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The memcopy only initializes part of the allocation, but mremap needs everything to be initialized.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh I see, thanks!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a comment explaining this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also seems worth saying that we assume here that zeroing the entire allocation is more efficient than zeroing just the parts not copied from the old allocation.

let old_size = match old_size_and_align {
Some((size, _align)) => size,
None => self.get_alloc_raw(alloc_id)?.size(),
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_const_eval/src/interpret/place.rs
Original file line number Diff line number Diff line change
Expand Up @@ -983,7 +983,7 @@ where
let Some((size, align)) = self.size_and_align_of(&meta, &layout)? else {
span_bug!(self.cur_span(), "cannot allocate space for `extern` type, size is not known")
};
let ptr = self.allocate_ptr(size, align, kind)?;
let ptr = self.allocate_ptr(size, align, kind, false)?;
interp_ok(self.ptr_with_meta_to_mplace(ptr.into(), meta, layout, /*unaligned*/ false))
}

Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_const_eval/src/interpret/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ pub(crate) fn create_static_alloc<'tcx>(
static_def_id: LocalDefId,
layout: TyAndLayout<'tcx>,
) -> InterpResult<'tcx, MPlaceTy<'tcx>> {
let alloc = Allocation::try_uninit(layout.size, layout.align.abi)?;
let alloc = Allocation::try_new(layout.size, layout.align.abi, false)?;
let alloc_id = ecx.tcx.reserve_and_set_static_alloc(static_def_id.into());
assert_eq!(ecx.machine.static_root_ids, None);
ecx.machine.static_root_ids = Some((alloc_id, static_def_id));
Expand Down
17 changes: 11 additions & 6 deletions compiler/rustc_middle/src/mir/interpret/allocation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,12 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
Allocation::from_bytes(slice, Align::ONE, Mutability::Not)
}

fn uninit_inner<R>(size: Size, align: Align, fail: impl FnOnce() -> R) -> Result<Self, R> {
fn new_inner<R>(
size: Size,
align: Align,
zero_init: bool,
fail: impl FnOnce() -> R,
) -> Result<Self, R> {
// We raise an error if we cannot create the allocation on the host.
// This results in an error that can happen non-deterministically, since the memory
// available to the compiler can change between runs. Normally queries are always
Expand All @@ -306,7 +311,7 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
Ok(Allocation {
bytes,
provenance: ProvenanceMap::new(),
init_mask: InitMask::new(size, false),
init_mask: InitMask::new(size, zero_init),
align,
mutability: Mutability::Mut,
extra: (),
Expand All @@ -315,8 +320,8 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {

/// Try to create an Allocation of `size` bytes, failing if there is not enough memory
/// available to the compiler to do so.
pub fn try_uninit<'tcx>(size: Size, align: Align) -> InterpResult<'tcx, Self> {
Self::uninit_inner(size, align, || {
pub fn try_new<'tcx>(size: Size, align: Align, zero_init: bool) -> InterpResult<'tcx, Self> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bools are a bit icky, as the call sites don't make it clear what the bool means. Maybe an enum with Uninit and Zeroed variants would be better? What do you think @RalfJung

Copy link
Author

@SpecificProtagonist SpecificProtagonist Jan 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I'm not happy with the unlabeled bools. An additional (very simple) possibility could be to use comments – this is used by existing code: mem_copy(ptr, new_ptr.into(), old_size.min(new_size), /*nonoverlapping*/ true).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah a dedicated enum makes sense.

Self::new_inner(size, align, zero_init, || {
ty::tls::with(|tcx| tcx.dcx().delayed_bug("exhausted memory during interpretation"));
InterpErrorKind::ResourceExhaustion(ResourceExhaustionInfo::MemoryExhausted)
})
Expand All @@ -328,8 +333,8 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
///
/// Example use case: To obtain an Allocation filled with specific data,
/// first call this function and then call write_scalar to fill in the right data.
pub fn uninit(size: Size, align: Align) -> Self {
match Self::uninit_inner(size, align, || {
pub fn new(size: Size, align: Align, zero_init: bool) -> Self {
match Self::new_inner(size, align, zero_init, || {
panic!(
"interpreter ran out of memory: cannot create allocation of {} bytes",
size.bytes()
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_middle/src/ty/vtable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ pub(super) fn vtable_allocation_provider<'tcx>(
let ptr_align = tcx.data_layout.pointer_align.abi;

let vtable_size = ptr_size * u64::try_from(vtable_entries.len()).unwrap();
let mut vtable = Allocation::uninit(vtable_size, ptr_align);
let mut vtable = Allocation::new(vtable_size, ptr_align, false);

// No need to do any alignment checks on the memory accesses below, because we know the
// allocation is correctly aligned as we created it above. Also we're only offsetting by
Expand Down
5 changes: 3 additions & 2 deletions compiler/rustc_smir/src/rustc_smir/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,8 @@ pub(crate) fn try_new_allocation<'tcx>(
.layout_of(rustc_middle::ty::TypingEnv::fully_monomorphized().as_query_input(ty))
.map_err(|e| e.stable(tables))?
.align;
let mut allocation = rustc_middle::mir::interpret::Allocation::uninit(size, align.abi);
let mut allocation =
rustc_middle::mir::interpret::Allocation::new(size, align.abi, false);
allocation
.write_scalar(&tables.tcx, alloc_range(Size::ZERO, size), scalar)
.map_err(|e| e.stable(tables))?;
Expand All @@ -69,7 +70,7 @@ pub(crate) fn try_new_allocation<'tcx>(
.layout_of(rustc_middle::ty::TypingEnv::fully_monomorphized().as_query_input(ty))
.map_err(|e| e.stable(tables))?;
let mut allocation =
rustc_middle::mir::interpret::Allocation::uninit(layout.size, layout.align.abi);
rustc_middle::mir::interpret::Allocation::new(layout.size, layout.align.abi, false);
allocation
.write_scalar(
&tables.tcx,
Expand Down
15 changes: 4 additions & 11 deletions src/tools/miri/src/shims/alloc.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
use std::iter;

use rustc_abi::{Align, Size};
use rustc_ast::expand::allocator::AllocatorKind;

Expand Down Expand Up @@ -83,15 +81,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
fn malloc(&mut self, size: u64, zero_init: bool) -> InterpResult<'tcx, Pointer> {
let this = self.eval_context_mut();
let align = this.malloc_align(size);
let ptr = this.allocate_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into())?;
if zero_init {
// We just allocated this, the access is definitely in-bounds and fits into our address space.
this.write_bytes_ptr(
ptr.into(),
iter::repeat(0u8).take(usize::try_from(size).unwrap()),
)
.unwrap();
}
let ptr = this.allocate_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into(), zero_init)?;
interp_ok(ptr.into())
}

Expand All @@ -115,6 +105,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
Align::from_bytes(align).unwrap(),
MiriMemoryKind::C.into(),
false
)?;
this.write_pointer(ptr, &memptr)?;
interp_ok(Scalar::from_i32(0))
Expand Down Expand Up @@ -147,6 +138,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(new_size),
new_align,
MiriMemoryKind::C.into(),
false
)?;
interp_ok(new_ptr.into())
}
Expand Down Expand Up @@ -187,6 +179,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
Align::from_bytes(align).unwrap(),
MiriMemoryKind::C.into(),
false
)?;
interp_ok(ptr.into())
}
Expand Down
11 changes: 3 additions & 8 deletions src/tools/miri/src/shims/foreign_items.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::collections::hash_map::Entry;
use std::io::Write;
use std::iter;
use std::path::Path;

use rustc_abi::{Align, AlignFromBytesError, Size};
Expand Down Expand Up @@ -509,6 +508,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
Align::from_bytes(align).unwrap(),
memory_kind.into(),
false
)?;

ecx.write_pointer(ptr, dest)
Expand Down Expand Up @@ -537,14 +537,8 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
Align::from_bytes(align).unwrap(),
MiriMemoryKind::Rust.into(),
true
)?;

// We just allocated this, the access is definitely in-bounds.
this.write_bytes_ptr(
ptr.into(),
iter::repeat(0u8).take(usize::try_from(size).unwrap()),
)
.unwrap();
this.write_pointer(ptr, dest)
});
}
Expand Down Expand Up @@ -604,6 +598,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(new_size),
align,
MiriMemoryKind::Rust.into(),
false
)?;
this.write_pointer(new_ptr, dest)
});
Expand Down
1 change: 1 addition & 0 deletions src/tools/miri/src/shims/unix/fs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1109,6 +1109,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
dirent_layout.align.abi,
MiriMemoryKind::Runtime.into(),
false
)?;
let entry: Pointer = entry.into();

Expand Down
10 changes: 1 addition & 9 deletions src/tools/miri/src/shims/unix/linux/mem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,16 +49,8 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(new_size),
align,
MiriMemoryKind::Mmap.into(),
true
)?;
if let Some(increase) = new_size.checked_sub(old_size) {
// We just allocated this, the access is definitely in-bounds and fits into our address space.
// mmap guarantees new mappings are zero-init.
this.write_bytes_ptr(
ptr.wrapping_offset(Size::from_bytes(old_size), this).into(),
std::iter::repeat(0u8).take(usize::try_from(increase).unwrap()),
)
.unwrap();
}

interp_ok(Scalar::from_pointer(ptr, this))
}
Expand Down
16 changes: 7 additions & 9 deletions src/tools/miri/src/shims/unix/mem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,15 +111,13 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
return interp_ok(this.eval_libc("MAP_FAILED"));
}

let ptr =
this.allocate_ptr(Size::from_bytes(map_length), align, MiriMemoryKind::Mmap.into())?;
// We just allocated this, the access is definitely in-bounds and fits into our address space.
// mmap guarantees new mappings are zero-init.
this.write_bytes_ptr(
ptr.into(),
std::iter::repeat(0u8).take(usize::try_from(map_length).unwrap()),
)
.unwrap();
let ptr = this.allocate_ptr(
Size::from_bytes(map_length),
align,
MiriMemoryKind::Mmap.into(),
// mmap guarantees new mappings are zero-init.
true
)?;

interp_ok(Scalar::from_pointer(ptr, this))
}
Expand Down
8 changes: 2 additions & 6 deletions src/tools/miri/src/shims/windows/foreign_items.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,13 +262,8 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
Align::from_bytes(align).unwrap(),
MiriMemoryKind::WinHeap.into(),
zero_init
)?;
if zero_init {
this.write_bytes_ptr(
ptr.into(),
iter::repeat(0u8).take(usize::try_from(size).unwrap()),
)?;
}
this.write_pointer(ptr, dest)?;
}
"HeapFree" => {
Expand Down Expand Up @@ -300,6 +295,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
Size::from_bytes(size),
Align::from_bytes(align).unwrap(),
MiriMemoryKind::WinHeap.into(),
false
)?;
this.write_pointer(new_ptr, dest)?;
}
Expand Down
Loading