From 185e3ef924a5ada1f15c074788022429bff30d03 Mon Sep 17 00:00:00 2001 From: SpecificProtagonist Date: Sat, 25 Jan 2025 00:54:51 +0100 Subject: [PATCH 1/2] miri: optimize zeroed alloc --- .../rustc_const_eval/src/interpret/memory.rs | 14 ++++++ .../src/mir/interpret/allocation.rs | 43 +++++++++++++++++++ src/tools/miri/src/shims/alloc.rs | 14 +++--- src/tools/miri/src/shims/foreign_items.rs | 9 +--- src/tools/miri/src/shims/unix/mem.rs | 13 +++--- .../miri/src/shims/windows/foreign_items.rs | 24 ++++++----- 6 files changed, 81 insertions(+), 36 deletions(-) diff --git a/compiler/rustc_const_eval/src/interpret/memory.rs b/compiler/rustc_const_eval/src/interpret/memory.rs index 2772c94d52b0b..e75bef5a302aa 100644 --- a/compiler/rustc_const_eval/src/interpret/memory.rs +++ b/compiler/rustc_const_eval/src/interpret/memory.rs @@ -250,6 +250,20 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> { self.insert_allocation(alloc, kind) } + pub fn allocate_zeroed_ptr( + &mut self, + size: Size, + align: Align, + kind: MemoryKind, + ) -> InterpResult<'tcx, Pointer> { + let alloc = if M::PANIC_ON_ALLOC_FAIL { + Allocation::zeroed(size, align) + } else { + Allocation::try_zeroed(size, align)? + }; + self.insert_allocation(alloc, kind) + } + pub fn insert_allocation( &mut self, alloc: Allocation, diff --git a/compiler/rustc_middle/src/mir/interpret/allocation.rs b/compiler/rustc_middle/src/mir/interpret/allocation.rs index 1b07846e0cf6c..37f90d887a068 100644 --- a/compiler/rustc_middle/src/mir/interpret/allocation.rs +++ b/compiler/rustc_middle/src/mir/interpret/allocation.rs @@ -340,6 +340,49 @@ impl Allocation { } } + fn zeroed_inner(size: Size, align: Align, fail: impl FnOnce() -> R) -> Result { + // 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 + // deterministic. However, we can be non-deterministic here because all uses of const + // evaluation (including ConstProp!) will make compilation fail (via hard error + // or ICE) upon encountering a `MemoryExhausted` error. + let bytes = Bytes::zeroed(size, align).ok_or_else(fail)?; + + Ok(Allocation { + bytes, + provenance: ProvenanceMap::new(), + init_mask: InitMask::new(size, true), + align, + mutability: Mutability::Mut, + extra: (), + }) + } + + /// Try to create an Allocation of `size` zero-initialized bytes, failing if there is not enough memory + /// available to the compiler to do so. + pub fn try_zeroed<'tcx>(size: Size, align: Align) -> InterpResult<'tcx, Self> { + Self::zeroed_inner(size, align, || { + ty::tls::with(|tcx| tcx.dcx().delayed_bug("exhausted memory during interpretation")); + InterpErrorKind::ResourceExhaustion(ResourceExhaustionInfo::MemoryExhausted) + }) + .into() + } + + /// Try to create an Allocation of `size` zero-initialized bytes, panics if there is not enough memory + /// available to the compiler to do so. + pub fn zeroed(size: Size, align: Align) -> Self { + match Self::zeroed_inner(size, align, || { + panic!( + "interpreter ran out of memory: cannot create allocation of {} bytes", + size.bytes() + ); + }) { + Ok(x) => x, + Err(x) => x, + } + } + /// Add the extra. pub fn with_extra(self, extra: Extra) -> Allocation { Allocation { diff --git a/src/tools/miri/src/shims/alloc.rs b/src/tools/miri/src/shims/alloc.rs index 25c0b52d0618b..b08ea388eeb3d 100644 --- a/src/tools/miri/src/shims/alloc.rs +++ b/src/tools/miri/src/shims/alloc.rs @@ -83,15 +83,11 @@ 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 = if zero_init { + this.allocate_zeroed_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into())? + } else { + this.allocate_ptr(Size::from_bytes(size), align, MiriMemoryKind::C.into())? + }; interp_ok(ptr.into()) } diff --git a/src/tools/miri/src/shims/foreign_items.rs b/src/tools/miri/src/shims/foreign_items.rs index 6c8ccc8398590..143cdbf0cdbb1 100644 --- a/src/tools/miri/src/shims/foreign_items.rs +++ b/src/tools/miri/src/shims/foreign_items.rs @@ -533,18 +533,11 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> { this.check_rustc_alloc_request(size, align)?; - let ptr = this.allocate_ptr( + let ptr = this.allocate_zeroed_ptr( Size::from_bytes(size), Align::from_bytes(align).unwrap(), MiriMemoryKind::Rust.into(), )?; - - // 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) }); } diff --git a/src/tools/miri/src/shims/unix/mem.rs b/src/tools/miri/src/shims/unix/mem.rs index 5531b944e17d1..0f6b12ac9d2c4 100644 --- a/src/tools/miri/src/shims/unix/mem.rs +++ b/src/tools/miri/src/shims/unix/mem.rs @@ -111,15 +111,12 @@ 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_zeroed_ptr( + Size::from_bytes(map_length), + align, + MiriMemoryKind::Mmap.into(), + )?; interp_ok(Scalar::from_pointer(ptr, this)) } diff --git a/src/tools/miri/src/shims/windows/foreign_items.rs b/src/tools/miri/src/shims/windows/foreign_items.rs index 0bf56c3d005fd..e5f63ef2ae938 100644 --- a/src/tools/miri/src/shims/windows/foreign_items.rs +++ b/src/tools/miri/src/shims/windows/foreign_items.rs @@ -258,17 +258,19 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> { // Alignment is twice the pointer size. // Source: let align = this.tcx.pointer_size().bytes().strict_mul(2); - let ptr = this.allocate_ptr( - Size::from_bytes(size), - Align::from_bytes(align).unwrap(), - MiriMemoryKind::WinHeap.into(), - )?; - if zero_init { - this.write_bytes_ptr( - ptr.into(), - iter::repeat(0u8).take(usize::try_from(size).unwrap()), - )?; - } + let ptr = if zero_init { + this.allocate_zeroed_ptr( + Size::from_bytes(size), + Align::from_bytes(align).unwrap(), + MiriMemoryKind::WinHeap.into(), + )? + } else { + this.allocate_ptr( + Size::from_bytes(size), + Align::from_bytes(align).unwrap(), + MiriMemoryKind::WinHeap.into(), + )? + }; this.write_pointer(ptr, dest)?; } "HeapFree" => { From bd28fafbc8b741ae7631a3b8543a2b1cddac35cf Mon Sep 17 00:00:00 2001 From: SpecificProtagonist Date: Sat, 25 Jan 2025 10:50:07 +0100 Subject: [PATCH 2/2] leftover import --- src/tools/miri/src/shims/alloc.rs | 2 -- src/tools/miri/src/shims/foreign_items.rs | 1 - 2 files changed, 3 deletions(-) diff --git a/src/tools/miri/src/shims/alloc.rs b/src/tools/miri/src/shims/alloc.rs index b08ea388eeb3d..fdead7a3aa2ca 100644 --- a/src/tools/miri/src/shims/alloc.rs +++ b/src/tools/miri/src/shims/alloc.rs @@ -1,5 +1,3 @@ -use std::iter; - use rustc_abi::{Align, Size}; use rustc_ast::expand::allocator::AllocatorKind; diff --git a/src/tools/miri/src/shims/foreign_items.rs b/src/tools/miri/src/shims/foreign_items.rs index 143cdbf0cdbb1..b78265147b0b4 100644 --- a/src/tools/miri/src/shims/foreign_items.rs +++ b/src/tools/miri/src/shims/foreign_items.rs @@ -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};