MMU-based epoch interruption

cranelift/codegen/meta/src/cdsl/instructions.rs

@@ -126,7 +126,7 @@ pub(crate) struct InstructionBuilder {
     operands_in: Option<Vec<Operand>>,
     operands_out: Option<Vec<Operand>>,
 
-    // See Instruction comments for the meaning of these fields.
+    // See InstructionContent comments for the meaning of these fields.
     is_terminator: bool,
     is_branch: bool,
     is_call: bool,

cranelift/codegen/meta/src/shared/instructions.rs

@@ -380,6 +380,41 @@ fn define_control_flow(
         .call()
         .branches(),
     );
+
+    ig.push(
+        Inst::new(
+            "dead_load_with_context",
+            r#"
+        Load 32 bits from memory at ``load_ptr`` while also keeping ``context``
+        in a fixed register and reserving a second as scratch space.
+
+        This is intended for implementing MMU-triggered jumps as in
+        epoch-interruption-via-mmu, where the load conditionally triggers a
+        segfault, which hands off control to a signal handler for further
+        action. The handler has access to ``context`` (typically the
+        ``VMContext``'s ``vm_store_context``) and can use the second
+        reserved register to store a temp value--like the original return
+        value--as needed on platforms where signal handlers cannot push stack
+        frames.
+
+        On x64, RDI holds ``context``, and R10 is used as scratch space.
+        "#,
+            &formats.binary,
+        )
+        .operands_in(vec![
+            Operand::new("load_ptr", iAddr).with_doc("memory location to load from"),
+            Operand::new("context", iAddr)
+                .with_doc("arbitrary address-sized context to pass to signal handler"),
+        ])
+        // Are we a call? stack_switch calls itself one "as it continues
+        // execution elsewhere". See reasoning at
+        // https://github.com/bytecodealliance/wasmtime/pull/9078#issuecomment-2273869774.
+        .call()
+        .can_load()
+        // Don't optimize me out just because I don't def anything. TODO: Can we use side_effects_idempotent()?
+        .other_side_effects(),
+        // If `load` is not can_trap(), this isn't either.
+    );
 }
 
 #[inline(never)]

cranelift/codegen/src/ir/instructions.rs

@@ -613,9 +613,13 @@ impl InstructionData {
             Self::Ternary {
                 opcode: Opcode::StackSwitch,
                 ..
+            }
+            | Self::Binary {
+                opcode: Opcode::DeadLoadWithContext,
+                ..
             } => {
-                // `StackSwitch` is not actually a call, but has the .call() side
-                // effect as it continues execution elsewhere.
+                // These instructions aren't actually calls, but they have the
+                // .call() side effect, as they continue execution elsewhere.
                 CallInfo::NotACall
             }
             _ => {

cranelift/codegen/src/isa/x64/inst.isle

@@ -52,7 +52,7 @@
        ;; =========================================
        ;; Stack manipulation.
 
-      ;; Emits a inline stack probe loop.
+      ;; Emits an inline stack probe loop.
       (StackProbeLoop (tmp WritableReg)
                       (frame_size u32)
                       (guard_size u32))
@@ -194,6 +194,10 @@
                     (offset i64)
                     (distance RelocDistance))
 
+       (DeadLoadWithContext (dst WritableGpr)
+                            (load_ptr Gpr)
+                            (context Gpr))
+
        ;; =========================================
        ;; Instructions pertaining to atomic memory accesses.
 

cranelift/codegen/src/isa/x64/inst/emit.rs

@@ -645,6 +645,20 @@ pub(crate) fn emit(
             sink.bind_label(resume, state.ctrl_plane_mut());
         }
 
+        Inst::DeadLoadWithContext { dst, load_ptr, .. } => {
+            let start = sink.cur_offset();
+
+            let load_ptr_addr = SyntheticAmode::real(Amode::imm_reg(0, **load_ptr));
+            // Since we're clobbering dst anyway to store the original return
+            // address, also use it as a destination for the dead load rather
+            // than sucking up another reg:
+            asm::inst::movq_rm::new(*dst, load_ptr_addr).emit(sink, info, state);
+
+            // Put the address of this instruction aside so we can later
+            // distinguish whether a segfault is its fault.
+            sink.add_epoch_check(start, sink.cur_offset());
+        }
+
         Inst::JmpKnown { dst } => uncond_jmp(sink, *dst),
 
         Inst::WinchJmpIf { cc, taken } => one_way_jmp(sink, *cc, *taken),

cranelift/codegen/src/isa/x64/inst/mod.rs

@@ -96,6 +96,7 @@ impl Inst {
             | Inst::Args { .. }
             | Inst::Rets { .. }
             | Inst::StackSwitchBasic { .. }
+            | Inst::DeadLoadWithContext { .. }
             | Inst::TrapIf { .. }
             | Inst::TrapIfAnd { .. }
             | Inst::TrapIfOr { .. }
@@ -156,7 +157,7 @@ impl Inst {
         Inst::External { inst }
     }
 
-    /// Writes the `simm64` immedaite into `dst`.
+    /// Writes the `simm64` immediate into `dst`.
     ///
     /// Note that if `dst_size` is less than 64-bits then the upper bits of
     /// `simm64` will be converted to zero.
@@ -671,6 +672,17 @@ impl PrettyPrint for Inst {
                 )
             }
 
+            Inst::DeadLoadWithContext {
+                dst,
+                load_ptr,
+                context,
+            } => {
+                let dst = pretty_print_reg(*dst.to_reg(), 8);
+                let load_ptr = pretty_print_reg(**load_ptr, 8);
+                let context = pretty_print_reg(**context, 8);
+                format!("dead_load_with_context {dst}, {load_ptr}, {context}")
+            }
+
             Inst::JmpKnown { dst } => {
                 let op = ljustify("jmp".to_string());
                 let dst = dst.to_string();
@@ -1045,6 +1057,26 @@ fn x64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_clobbers(clobbers);
         }
 
+        Inst::DeadLoadWithContext {
+            dst,
+            load_ptr,
+            context,
+        } => {
+            // load_ptr is an input param.
+            collector.reg_use(load_ptr);
+            // Demand context (vmctx) go into RDI.
+            // TODO: Do I still have to move it, or will regalloc make sure it's there?
+            collector.reg_fixed_use(context, regs::rdi());
+            // Reserve r10 as a place for the signal handler to stow the return
+            // address (which we're overwriting with that of the epoch-ending
+            // stub). Picking r10 because it's caller-saved but otherwise
+            // arbitrarily.
+            //
+            // Also def it so we can use it as the destination of the dead load
+            // rather than consuming another arbitrary reg.
+            collector.reg_fixed_def(dst, regs::r10());
+        }
+
         Inst::ReturnCallKnown { info } => {
             let ReturnCallInfo {
                 dest, uses, tmp, ..

cranelift/codegen/src/isa/x64/lower.isle

@@ -3578,6 +3578,18 @@
             (in_payload0 Gpr (put_in_gpr in_payload0)))
         (x64_stack_switch_basic store_context_ptr load_context_ptr in_payload0)))
 
+;;;; Rules for `dead_load_with_context` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(rule (lower (dead_load_with_context load_ptr context))
+      (let ((load_ptr Gpr (put_in_gpr load_ptr))
+            (context Gpr (put_in_gpr context))
+            (dst WritableGpr (temp_writable_gpr))
+            (_ Unit (emit_side_effect (SideEffectNoResult.Inst
+                  (MInst.DeadLoadWithContext dst
+                                             load_ptr
+                                             context)))))
+       (output_none)))
+
 ;;;; Rules for `get_{frame,stack}_pointer` and `get_return_address` ;;;;;;;;;;;;
 
 (rule (lower (get_frame_pointer))

cranelift/codegen/src/isa/x64/pcc.rs

@@ -209,6 +209,8 @@ pub(crate) fn check(
 
         Inst::StackSwitchBasic { .. } => Err(PccError::UnimplementedInst),
 
+        Inst::DeadLoadWithContext { .. } => Err(PccError::UnimplementedInst),
+
         Inst::LabelAddress { .. } => Err(PccError::UnimplementedInst),
 
         Inst::SequencePoint { .. } => Ok(()),

cranelift/codegen/src/isle_prelude.rs

@@ -638,6 +638,11 @@ macro_rules! isle_common_prelude_methods {
             MemFlags::trusted()
         }
 
+        #[inline]
+        fn mem_flags_aligned_read_only(&mut self) -> MemFlags {
+            MemFlags::new().with_aligned().with_readonly()
+        }
+
         #[inline]
         fn little_or_native_endian(&mut self, flags: MemFlags) -> Option<MemFlags> {
             match flags.explicit_endianness() {

cranelift/codegen/src/machinst/buffer.rs

@@ -252,6 +252,8 @@ pub struct MachBuffer<I: VCodeInst> {
     call_sites: SmallVec<[MachCallSite; 16]>,
     /// Any patchable call site locations.
     patchable_call_sites: SmallVec<[MachPatchableCallSite; 16]>,
+    /// Any locations which do an MMU-based check for the end of an epoch.
+    epoch_checks: SmallVec<[Range<CodeOffset>; 16]>,
     /// Any exception-handler records referred to at call sites.
     exception_handlers: SmallVec<[MachExceptionHandler; 16]>,
     /// Any source location mappings referring to this code.
@@ -343,6 +345,7 @@ impl MachBufferFinalized<Stencil> {
             traps: self.traps,
             call_sites: self.call_sites,
             patchable_call_sites: self.patchable_call_sites,
+            epoch_checks: self.epoch_checks,
             exception_handlers: self.exception_handlers,
             srclocs: self
                 .srclocs
@@ -380,6 +383,8 @@ pub struct MachBufferFinalized<T: CompilePhase> {
     pub(crate) call_sites: SmallVec<[MachCallSite; 16]>,
     /// Any patchable call site locations refering to this code.
     pub(crate) patchable_call_sites: SmallVec<[MachPatchableCallSite; 16]>,
+    /// Any locations which do an MMU-based check for the end of an epoch.
+    pub epoch_checks: SmallVec<[Range<CodeOffset>; 16]>,
     /// Any exception-handler records referred to at call sites.
     pub(crate) exception_handlers: SmallVec<[FinalizedMachExceptionHandler; 16]>,
     /// Any source location mappings referring to this code.
@@ -480,6 +485,7 @@ impl<I: VCodeInst> MachBuffer<I> {
             traps: SmallVec::new(),
             call_sites: SmallVec::new(),
             patchable_call_sites: SmallVec::new(),
+            epoch_checks: SmallVec::new(),
             exception_handlers: SmallVec::new(),
             srclocs: SmallVec::new(),
             debug_tags: vec![],
@@ -1581,6 +1587,7 @@ impl<I: VCodeInst> MachBuffer<I> {
             traps: self.traps,
             call_sites: self.call_sites,
             patchable_call_sites: self.patchable_call_sites,
+            epoch_checks: self.epoch_checks,
             exception_handlers: finalized_exception_handlers,
             srclocs,
             debug_tags: self.debug_tags,
@@ -1685,7 +1692,7 @@ impl<I: VCodeInst> MachBuffer<I> {
         });
     }
 
-    /// Add a patchable call record at the current offset The actual
+    /// Add a patchable call record at the current offset. The actual
     /// call is expected to have been emitted; the VCodeInst trait
     /// specifies how to NOP it out, and we carry that information to
     /// the finalized Machbuffer.
@@ -1696,6 +1703,16 @@ impl<I: VCodeInst> MachBuffer<I> {
         });
     }
 
+    /// Record that an MMU-based epoch interruption check occurs at the current
+    /// offset. A signal handler may use these annotations to distinguish that a
+    /// segfault is actually an epoch interruption in disguise. The
+    /// DeadLoadWithContext instruction is assumed to have already been emitted.
+    /// `start` is the offset of the emitted instruction, and `end` is the
+    /// offset of the immediately following instruction.
+    pub fn add_epoch_check(&mut self, start: CodeOffset, end: CodeOffset) {
+        self.epoch_checks.push(start..end);
+    }
+
     /// Add an unwind record at the current offset.
     pub fn add_unwind(&mut self, unwind: UnwindInst) {
         self.unwind_info.push((self.cur_offset(), unwind));

cranelift/codegen/src/prelude.isle

@@ -304,10 +304,15 @@
 
 ;; `MemFlags::trusted`
 (spec (mem_flags_trusted)
-      (provide (= result #x0003)))
+      (provide (= result #x0003))) ;; Shouldn't this be 0001?
 (decl pure mem_flags_trusted () MemFlags)
 (extern constructor mem_flags_trusted mem_flags_trusted)
 
+(spec (mem_flags_aligned_read_only)
+      (provide (= result #x0003)))
+(decl pure mem_flags_aligned_read_only () MemFlags)
+(extern constructor mem_flags_aligned_read_only mem_flags_aligned_read_only)
+
 ;; Determine if flags specify little- or native-endian.
 (decl little_or_native_endian (MemFlags) MemFlags)
 (extern extractor little_or_native_endian little_or_native_endian)

cranelift/docs/ir.md

@@ -644,7 +644,7 @@ slot on the stack for its entire live range. Since the live range of an SSA
 value can be quite large, it is sometimes beneficial to split the live range
 into smaller parts.
 
-A live range is split by creating new SSA values that are copies or the
+A live range is split by creating new SSA values that are copies of the
 original value or each other. The copies are created by inserting `copy`,
 `spill`, or `fill` instructions, depending on whether the values
 are assigned to registers or stack slots.

cranelift/interpreter/src/step.rs

@@ -1316,6 +1316,7 @@ where
         Opcode::X86Pmaddubsw => unimplemented!("X86Pmaddubsw"),
         Opcode::X86Cvtt2dq => unimplemented!("X86Cvtt2dq"),
         Opcode::StackSwitch => unimplemented!("StackSwitch"),
+        Opcode::DeadLoadWithContext => unimplemented!("DeadLoadWithContext"),
 
         Opcode::TryCall => unimplemented!("TryCall"),
         Opcode::TryCallIndirect => unimplemented!("TryCallIndirect"),

crates/cli-flags/src/lib.rs

@@ -298,6 +298,10 @@ wasmtime_option_group! {
         /// Yield when a global epoch counter changes, allowing for async
         /// operation without blocking the executor.
         pub epoch_interruption: Option<bool>,
+        /// Use MMU tricks to speed epoch deadline checks.
+        /// TODO: Document whether this should be used mutually exclusively with
+        /// epoch_interruption.
+        pub epoch_interruption_via_mmu: Option<bool>,
         /// Maximum stack size, in bytes, that wasm is allowed to consume before a
         /// stack overflow is reported.
         pub max_wasm_stack: Option<usize>,
@@ -830,6 +834,9 @@ impl CommonOptions {
         if let Some(enable) = self.wasm.epoch_interruption {
             config.epoch_interruption(enable);
         }
+        if let Some(enable) = self.wasm.epoch_interruption_via_mmu {
+            config.epoch_interruption_via_mmu(enable);
+        }
         if let Some(enable) = self.debug.address_map {
             config.generate_address_map(enable);
         }

crates/cranelift/src/compiler.rs

@@ -35,10 +35,11 @@ use wasmparser::{FuncValidatorAllocations, FunctionBody};
 use wasmtime_environ::obj::{ELF_WASMTIME_EXCEPTIONS, ELF_WASMTIME_FRAMES};
 use wasmtime_environ::{
     Abi, AddressMapSection, BuiltinFunctionIndex, CacheStore, CompileError, CompiledFunctionBody,
-    DefinedFuncIndex, FlagValue, FrameInstPos, FrameStackShape, FrameStateSlotBuilder,
-    FrameTableBuilder, FuncKey, FunctionBodyData, FunctionLoc, HostCall, InliningCompiler,
-    ModuleTranslation, ModuleTypesBuilder, PtrSize, StackMapSection, StaticModuleIndex,
-    TrapEncodingBuilder, TrapSentinel, TripleExt, Tunables, WasmFuncType, WasmValType,
+    DefinedFuncIndex, EpochCheckSection, FlagValue, FrameInstPos, FrameStackShape,
+    FrameStateSlotBuilder, FrameTableBuilder, FuncKey, FunctionBodyData, FunctionLoc, HostCall,
+    InliningCompiler, ModuleTranslation, ModuleTypesBuilder, PtrSize, StackMapSection,
+    StaticModuleIndex, TrapEncodingBuilder, TrapSentinel, TripleExt, Tunables, WasmFuncType,
+    WasmValType,
 };
 use wasmtime_unwinder::ExceptionTableBuilder;
 
@@ -468,6 +469,7 @@ impl wasmtime_environ::Compiler for Compiler {
         let mut stack_maps = StackMapSection::default();
         let mut exception_tables = ExceptionTableBuilder::default();
         let mut frame_tables = FrameTableBuilder::default();
+        let mut epoch_checks = EpochCheckSection::default();
 
         let funcs = funcs
             .iter()
@@ -536,6 +538,9 @@ impl wasmtime_environ::Compiler for Compiler {
                 )?;
                 nop_units.get_or_insert_with(|| func.buffer.nop_units.clone());
             }
+            if self.tunables.epoch_interruption_via_mmu {
+                epoch_checks.push(range.clone(), &func.buffer.epoch_checks);
+            }
             builder.append_padding(self.linkopts.padding_between_functions);
 
             let info = FunctionLoc {
@@ -582,6 +587,9 @@ impl wasmtime_environ::Compiler for Compiler {
         }
         stack_maps.append_to(obj);
         traps.append_to(obj);
+        if self.tunables.epoch_interruption_via_mmu {
+            epoch_checks.append_to(obj);
+        }
 
         let exception_section = obj.add_section(
             obj.segment_name(StandardSegment::Data).to_vec(),