From 14801af08fd270940a4812083a0271b179a8f1bd Mon Sep 17 00:00:00 2001 From: Ralf Jung Date: Mon, 27 May 2024 23:28:32 +0200 Subject: [PATCH] miri: avoid making a full copy of all new allocations --- .../src/const_eval/dummy_machine.rs | 2 +- .../src/const_eval/machine.rs | 2 +- .../src/interpret/eval_context.rs | 2 +- .../rustc_const_eval/src/interpret/machine.rs | 58 +++++++++++++------ .../rustc_const_eval/src/interpret/memory.rs | 14 +++-- .../src/mir/interpret/allocation.rs | 43 +++++++------- src/tools/miri/src/concurrency/thread.rs | 13 +++-- src/tools/miri/src/machine.rs | 40 +++++-------- 8 files changed, 96 insertions(+), 78 deletions(-) diff --git a/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs b/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs index 62979af8a601d..9a98677a84482 100644 --- a/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs +++ b/compiler/rustc_const_eval/src/const_eval/dummy_machine.rs @@ -174,7 +174,7 @@ impl<'tcx> interpret::Machine<'tcx> for DummyMachine { unimplemented!() } - fn init_frame_extra( + fn init_frame( _ecx: &mut InterpCx<'tcx, Self>, _frame: interpret::Frame<'tcx, Self::Provenance>, ) -> interpret::InterpResult<'tcx, interpret::Frame<'tcx, Self::Provenance, Self::FrameExtra>> diff --git a/compiler/rustc_const_eval/src/const_eval/machine.rs b/compiler/rustc_const_eval/src/const_eval/machine.rs index 950e7c12da2e9..2c0cbc9249352 100644 --- a/compiler/rustc_const_eval/src/const_eval/machine.rs +++ b/compiler/rustc_const_eval/src/const_eval/machine.rs @@ -655,7 +655,7 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeInterpreter<'tcx> { } #[inline(always)] - fn init_frame_extra( + fn init_frame( ecx: &mut InterpCx<'tcx, Self>, frame: Frame<'tcx>, ) -> InterpResult<'tcx, Frame<'tcx>> { diff --git a/compiler/rustc_const_eval/src/interpret/eval_context.rs b/compiler/rustc_const_eval/src/interpret/eval_context.rs index 5cd50a928704d..7c2100fcbe38a 100644 --- a/compiler/rustc_const_eval/src/interpret/eval_context.rs +++ b/compiler/rustc_const_eval/src/interpret/eval_context.rs @@ -819,7 +819,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> { tracing_span: SpanGuard::new(), extra: (), }; - let frame = M::init_frame_extra(self, pre_frame)?; + let frame = M::init_frame(self, pre_frame)?; self.stack_mut().push(frame); // Make sure all the constants required by this frame evaluate successfully (post-monomorphization check). diff --git a/compiler/rustc_const_eval/src/interpret/machine.rs b/compiler/rustc_const_eval/src/interpret/machine.rs index 5713e7bd82b3f..4ae0aca5a0c09 100644 --- a/compiler/rustc_const_eval/src/interpret/machine.rs +++ b/compiler/rustc_const_eval/src/interpret/machine.rs @@ -329,27 +329,41 @@ pub trait Machine<'tcx>: Sized { ptr: Pointer, ) -> Option<(AllocId, Size, Self::ProvenanceExtra)>; - /// Called to adjust allocations to the Provenance and AllocExtra of this machine. + /// Called to adjust global allocations to the Provenance and AllocExtra of this machine. /// /// If `alloc` contains pointers, then they are all pointing to globals. /// - /// The way we construct allocations is to always first construct it without extra and then add - /// the extra. This keeps uniform code paths for handling both allocations created by CTFE for - /// globals, and allocations created by Miri during evaluation. - /// - /// `kind` is the kind of the allocation being adjusted; it can be `None` when - /// it's a global and `GLOBAL_KIND` is `None`. - /// /// This should avoid copying if no work has to be done! If this returns an owned /// allocation (because a copy had to be done to adjust things), machine memory will /// cache the result. (This relies on `AllocMap::get_or` being able to add the /// owned allocation to the map even when the map is shared.) - fn adjust_allocation<'b>( + fn adjust_global_allocation<'b>( ecx: &InterpCx<'tcx, Self>, id: AllocId, - alloc: Cow<'b, Allocation>, - kind: Option>, - ) -> InterpResult<'tcx, Cow<'b, Allocation>>; + alloc: &'b Allocation, + ) -> InterpResult<'tcx, Cow<'b, Allocation>> + { + // The default implementation does a copy; CTFE machines have a more efficient implementation + // based on their particular choice for `Provenance`, `AllocExtra`, and `Bytes`. + let kind = Self::GLOBAL_KIND + .expect("if GLOBAL_KIND is None, adjust_global_allocation must be overwritten"); + let alloc = alloc.adjust_from_tcx(&ecx.tcx, |ptr| ecx.global_root_pointer(ptr))?; + let extra = + Self::init_alloc_extra(ecx, id, MemoryKind::Machine(kind), alloc.size(), alloc.align)?; + Ok(Cow::Owned(alloc.with_extra(extra))) + } + + /// Initialize the extra state of an allocation. + /// + /// This is guaranteed to be called exactly once on all allocations that are accessed by the + /// program. + fn init_alloc_extra( + ecx: &InterpCx<'tcx, Self>, + id: AllocId, + kind: MemoryKind, + size: Size, + align: Align, + ) -> InterpResult<'tcx, Self::AllocExtra>; /// Return a "root" pointer for the given allocation: the one that is used for direct /// accesses to this static/const/fn allocation, or the one returned from the heap allocator. @@ -473,7 +487,7 @@ pub trait Machine<'tcx>: Sized { } /// Called immediately before a new stack frame gets pushed. - fn init_frame_extra( + fn init_frame( ecx: &mut InterpCx<'tcx, Self>, frame: Frame<'tcx, Self::Provenance>, ) -> InterpResult<'tcx, Frame<'tcx, Self::Provenance, Self::FrameExtra>>; @@ -590,13 +604,23 @@ pub macro compile_time_machine(<$tcx: lifetime>) { } #[inline(always)] - fn adjust_allocation<'b>( + fn adjust_global_allocation<'b>( _ecx: &InterpCx<$tcx, Self>, _id: AllocId, - alloc: Cow<'b, Allocation>, - _kind: Option>, + alloc: &'b Allocation, ) -> InterpResult<$tcx, Cow<'b, Allocation>> { - Ok(alloc) + // Overwrite default implementation: no need to adjust anything. + Ok(Cow::Borrowed(alloc)) + } + + fn init_alloc_extra( + _ecx: &InterpCx<$tcx, Self>, + _id: AllocId, + _kind: MemoryKind, + _size: Size, + _align: Align, + ) -> InterpResult<$tcx, Self::AllocExtra> { + Ok(()) } fn extern_static_pointer( diff --git a/compiler/rustc_const_eval/src/interpret/memory.rs b/compiler/rustc_const_eval/src/interpret/memory.rs index 40bbfaa92c6fd..521f28b7123ea 100644 --- a/compiler/rustc_const_eval/src/interpret/memory.rs +++ b/compiler/rustc_const_eval/src/interpret/memory.rs @@ -239,7 +239,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> { pub fn allocate_raw_ptr( &mut self, - alloc: Allocation, + alloc: Allocation, kind: MemoryKind, ) -> InterpResult<'tcx, Pointer> { let id = self.tcx.reserve_alloc_id(); @@ -248,8 +248,11 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> { M::GLOBAL_KIND.map(MemoryKind::Machine), "dynamically allocating global memory" ); - let alloc = M::adjust_allocation(self, id, Cow::Owned(alloc), Some(kind))?; - self.memory.alloc_map.insert(id, (kind, alloc.into_owned())); + // We have set things up so we don't need to call `adjust_from_tcx` here, + // so we avoid copying the entire allocation contents. + let extra = M::init_alloc_extra(self, id, kind, alloc.size(), alloc.align)?; + let alloc = alloc.with_extra(extra); + self.memory.alloc_map.insert(id, (kind, alloc)); M::adjust_alloc_root_pointer(self, Pointer::from(id), Some(kind)) } @@ -583,11 +586,10 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> { }; M::before_access_global(self.tcx, &self.machine, id, alloc, def_id, is_write)?; // We got tcx memory. Let the machine initialize its "extra" stuff. - M::adjust_allocation( + M::adjust_global_allocation( self, id, // always use the ID we got as input, not the "hidden" one. - Cow::Borrowed(alloc.inner()), - M::GLOBAL_KIND.map(MemoryKind::Machine), + alloc.inner(), ) } diff --git a/compiler/rustc_middle/src/mir/interpret/allocation.rs b/compiler/rustc_middle/src/mir/interpret/allocation.rs index 791e87735f40b..b0f8a047b82f9 100644 --- a/compiler/rustc_middle/src/mir/interpret/allocation.rs +++ b/compiler/rustc_middle/src/mir/interpret/allocation.rs @@ -266,19 +266,6 @@ impl AllocRange { // The constructors are all without extra; the extra gets added by a machine hook later. impl Allocation { - /// Creates an allocation from an existing `Bytes` value - this is needed for miri FFI support - pub fn from_raw_bytes(bytes: Bytes, align: Align, mutability: Mutability) -> Self { - let size = Size::from_bytes(bytes.len()); - Self { - bytes, - provenance: ProvenanceMap::new(), - init_mask: InitMask::new(size, true), - align, - mutability, - extra: (), - } - } - /// Creates an allocation initialized by the given bytes pub fn from_bytes<'a>( slice: impl Into>, @@ -342,18 +329,30 @@ impl Allocation { Err(x) => x, } } + + /// Add the extra. + pub fn with_extra(self, extra: Extra) -> Allocation { + Allocation { + bytes: self.bytes, + provenance: self.provenance, + init_mask: self.init_mask, + align: self.align, + mutability: self.mutability, + extra, + } + } } impl Allocation { /// Adjust allocation from the ones in `tcx` to a custom Machine instance - /// with a different `Provenance`, `Extra` and `Byte` type. - pub fn adjust_from_tcx( - self, + /// with a different `Provenance` and `Byte` type. + pub fn adjust_from_tcx( + &self, cx: &impl HasDataLayout, - extra: Extra, mut adjust_ptr: impl FnMut(Pointer) -> Result, Err>, - ) -> Result, Err> { - let mut bytes = self.bytes; + ) -> Result, Err> { + // Copy the data. + let mut bytes = Bytes::from_bytes(Cow::Borrowed(&*self.bytes), self.align); // Adjust provenance of pointers stored in this allocation. let mut new_provenance = Vec::with_capacity(self.provenance.ptrs().len()); let ptr_size = cx.data_layout().pointer_size.bytes_usize(); @@ -369,12 +368,12 @@ impl Allocation { } // Create allocation. Ok(Allocation { - bytes: AllocBytes::from_bytes(Cow::Owned(Vec::from(bytes)), self.align), + bytes, provenance: ProvenanceMap::from_presorted_ptrs(new_provenance), - init_mask: self.init_mask, + init_mask: self.init_mask.clone(), align: self.align, mutability: self.mutability, - extra, + extra: self.extra, }) } } diff --git a/src/tools/miri/src/concurrency/thread.rs b/src/tools/miri/src/concurrency/thread.rs index 8e8e1182bfffe..06af72dadbd38 100644 --- a/src/tools/miri/src/concurrency/thread.rs +++ b/src/tools/miri/src/concurrency/thread.rs @@ -862,14 +862,15 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> { if tcx.is_foreign_item(def_id) { throw_unsup_format!("foreign thread-local statics are not supported"); } - let allocation = this.ctfe_query(|tcx| tcx.eval_static_initializer(def_id))?; - let mut allocation = allocation.inner().clone(); + let alloc = this.ctfe_query(|tcx| tcx.eval_static_initializer(def_id))?; + // We make a full copy of this allocation. + let mut alloc = alloc.inner().adjust_from_tcx(&this.tcx, |ptr| this.global_root_pointer(ptr))?; // This allocation will be deallocated when the thread dies, so it is not in read-only memory. - allocation.mutability = Mutability::Mut; + alloc.mutability = Mutability::Mut; // Create a fresh allocation with this content. - let new_alloc = this.allocate_raw_ptr(allocation, MiriMemoryKind::Tls.into())?; - this.machine.threads.set_thread_local_alloc(def_id, new_alloc); - Ok(new_alloc) + let ptr = this.allocate_raw_ptr(alloc, MiriMemoryKind::Tls.into())?; + this.machine.threads.set_thread_local_alloc(def_id, ptr); + Ok(ptr) } } diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs index fc1a655c1232f..280edfa585b29 100644 --- a/src/tools/miri/src/machine.rs +++ b/src/tools/miri/src/machine.rs @@ -1,7 +1,6 @@ //! Global machine state as well as implementation of the interpreter engine //! `Machine` trait. -use std::borrow::Cow; use std::cell::RefCell; use std::collections::hash_map::Entry; use std::fmt; @@ -1086,35 +1085,34 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> { } } - fn adjust_allocation<'b>( + fn init_alloc_extra( ecx: &MiriInterpCx<'tcx>, id: AllocId, - alloc: Cow<'b, Allocation>, - kind: Option, - ) -> InterpResult<'tcx, Cow<'b, Allocation>> + kind: MemoryKind, + size: Size, + align: Align, + ) -> InterpResult<'tcx, Self::AllocExtra> { - let kind = kind.expect("we set our STATIC_KIND so this cannot be None"); if ecx.machine.tracked_alloc_ids.contains(&id) { ecx.emit_diagnostic(NonHaltingDiagnostic::CreatedAlloc( id, - alloc.size(), - alloc.align, + size, + align, kind, )); } - let alloc = alloc.into_owned(); let borrow_tracker = ecx .machine .borrow_tracker .as_ref() - .map(|bt| bt.borrow_mut().new_allocation(id, alloc.size(), kind, &ecx.machine)); + .map(|bt| bt.borrow_mut().new_allocation(id, size, kind, &ecx.machine)); let race_alloc = ecx.machine.data_race.as_ref().map(|data_race| { data_race::AllocState::new_allocation( data_race, &ecx.machine.threads, - alloc.size(), + size, kind, ecx.machine.current_span(), ) @@ -1130,17 +1128,6 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> { Some(ecx.generate_stacktrace()) }; - let alloc: Allocation = alloc.adjust_from_tcx( - &ecx.tcx, - AllocExtra { - borrow_tracker, - data_race: race_alloc, - weak_memory: buffer_alloc, - backtrace, - }, - |ptr| ecx.global_root_pointer(ptr), - )?; - if matches!(kind, MemoryKind::Machine(kind) if kind.should_save_allocation_span()) { ecx.machine .allocation_spans @@ -1148,7 +1135,12 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> { .insert(id, (ecx.machine.current_span(), None)); } - Ok(Cow::Owned(alloc)) + Ok(AllocExtra { + borrow_tracker, + data_race: race_alloc, + weak_memory: buffer_alloc, + backtrace, + }) } fn adjust_alloc_root_pointer( @@ -1357,7 +1349,7 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> { } #[inline(always)] - fn init_frame_extra( + fn init_frame( ecx: &mut InterpCx<'tcx, Self>, frame: Frame<'tcx, Provenance>, ) -> InterpResult<'tcx, Frame<'tcx, Provenance, FrameExtra<'tcx>>> {