MMU-based epoch interruption

Cargo.toml

@@ -532,6 +532,11 @@ default = [
 #
 # These features are off-by-default but may optionally be enabled.
 all-arch = ["wasmtime/all-arch"]
+# Internal feature ensuring async runtime support is compiled in. Features that
+# require async (like `run`, `wizer`, and `serve`) forward to this, giving
+# `build.rs` a single, stable feature to check rather than chasing which
+# top-level features happen to pull in async support.
+async = ["wasmtime-cli-flags/async"]
 winch = ["wasmtime/winch"]
 wmemcheck = ["wasmtime/wmemcheck"]
 trace-log = ["wasmtime/trace-log"]
@@ -604,7 +609,7 @@ serve = [
   "dep:http-body-util",
   "dep:http",
   "dep:pin-project-lite",
-  "wasmtime-cli-flags/async",
+  "async",
   "wasmtime-wasi-http?/p2",
 ]
 explore = ["dep:wasmtime-explorer", "dep:tempfile"]
@@ -616,7 +621,7 @@ run = [
   "wasmtime/runtime",
   "wasmtime/wave",
   "dep:tokio",
-  "wasmtime-cli-flags/async",
+  "async",
   "wasmtime-wasi-http?/p2",
   "dep:wasmtime-debugger",
 ]
@@ -639,6 +644,10 @@ wizer = [
   "dep:wasmtime-wasi",
   "dep:tokio",
   "wasmtime/wave",
+  # dep:wasmtime-wasi transitively enables wasmtime-cli-flags/async, but making
+  # it explicit here makes the expression determining the has_mmu_interruption
+  # cfg flag's value clear.
+  "async",
 ]
 
 [[test]]

build.rs

@@ -14,6 +14,21 @@ fn main() {
         }
         _ => {}
     }
+
+    // Mirror the `has_mmu_interruption` cfg in `crates/wasmtime/build.rs`.
+    //
+    // We can omit the `std` check here, because `wasmtime` is always built with
+    // `std` from this crate.
+    println!("cargo:rustc-check-cfg=cfg(has_mmu_interruption)");
+    let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap();
+    let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
+    if target_os == "linux"
+        && (target_arch == "x86_64" || target_arch == "aarch64")
+        && cfg!(feature = "cranelift")
+        && cfg!(feature = "async")
+    {
+        println!("cargo:rustc-cfg=has_mmu_interruption");
+    }
 }
 
 fn set_commit_info_for_rustc() {

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

@@ -131,7 +131,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,43 @@ fn define_control_flow(
         .call()
         .branches(),
     );
+
+    ig.push(
+        Inst::new(
+            "dead_load_with_context",
+            r#"
+        Load a pointer-sized value 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
+        `mmu-interruption`, 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.
+
+        Which registers these are is ISA-specific; see each backend's
+        ``get_operands`` for the choices and the reasoning behind them.
+        "#,
+            &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

@@ -616,9 +616,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/aarch64/inst.isle

@@ -1013,7 +1013,13 @@
        ;; means that the internal codegen can't use these registers.
        (StackProbeLoop (start WritableReg)
                        (end Reg)
-                       (step Imm12))))
+                       (step Imm12))
+
+       ;; A load whose result is discarded; `context` is pinned to x0 and `dst`
+       ;; to x9.
+       (DeadLoadWithContext (dst WritableReg)
+                            (load_ptr Reg)
+                            (context Reg))))
 
 (spec (MInst.AluRRImmLogic alu_op size rd rn imml)
       (provide

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

@@ -3633,6 +3633,27 @@ impl MachInstEmit for Inst {
                 .emit(sink, emit_info, state);
                 sink.bind_label(loop_end, &mut state.ctrl_plane);
             }
+
+            &Inst::DeadLoadWithContext { dst, load_ptr, .. } => {
+                let start = sink.cur_offset();
+
+                // Emit `ldr dst, [load_ptr]`. Reuse the `dst` address as the
+                // destination of the dead load, since we are clobbering it
+                // anyway.
+                Inst::ULoad64 {
+                    rd: dst,
+                    mem: AMode::UnsignedOffset {
+                        rn: load_ptr,
+                        uimm12: UImm12Scaled::zero(I64),
+                    },
+                    flags: MemFlagsData::trusted(),
+                }
+                .emit(sink, emit_info, state);
+
+                // Mark the address of this instruction as part of mmu
+                // interrupt.
+                sink.add_mmu_interrupt_check(start, sink.cur_offset());
+            }
         }
 
         let end_off = sink.cur_offset();

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

@@ -926,6 +926,25 @@ fn aarch64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_early_def(start);
             collector.reg_use(end);
         }
+        Inst::DeadLoadWithContext {
+            dst,
+            load_ptr,
+            context,
+        } => {
+            // `load_ptr` is an ordinary input.
+            collector.reg_use(load_ptr);
+            // Demand `context` (the vmctx) go into x0, where the signal
+            // handler can find it and hand it straight to
+            // `task_switch_trampoline` as its first argument.
+            collector.reg_fixed_use(context, regs::xreg(0));
+            // Reserve x9 as a place for the signal handler to put the address
+            // at which to resume once the task switch is done. x9 is caller
+            // saved and has no special role.
+            //
+            // Define 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::xreg(9));
+        }
     }
 }
 
@@ -2929,6 +2948,16 @@ impl Inst {
                 let step = step.pretty_print(0);
                 format!("stack_probe_loop {start}, {end}, {step}")
             }
+            &Inst::DeadLoadWithContext {
+                dst,
+                load_ptr,
+                context,
+            } => {
+                let dst = pretty_print_reg(dst.to_reg());
+                let load_ptr = pretty_print_reg(load_ptr);
+                let context = pretty_print_reg(context);
+                format!("dead_load_with_context {dst}, {load_ptr}, {context}")
+            }
         }
     }
 }

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

@@ -2440,6 +2440,18 @@
 (rule (lower (symbol_value _ (symbol_value_data extname dist offset)))
       (load_ext_name (box_external_name extname) offset dist))
 
+;;;; Rules for `dead_load_with_context` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(rule (lower (dead_load_with_context load_ptr context))
+      (let ((load_ptr Reg (put_in_reg load_ptr))
+            (context Reg (put_in_reg context))
+            (dst WritableReg (temp_writable_reg $I64))
+            (_ 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/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_mmu_interrupt_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 { .. }
@@ -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();
@@ -1051,6 +1063,30 @@ 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.
+            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 and not used for arg
+            // passing in Linux/x64. It is used as "a static chain pointer
+            // in case of nested functions" according to SystemV, but that's
+            // inapplicable to compiled Wasm code. It is also used to store the
+            // function stack limit in Cranelift, but the stack-limit check is
+            // over by the time we need r10, in the case of the use of this
+            // instruction for MMU-based epoch interruption.
+            //
+            // 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

@@ -3582,6 +3582,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/isle_prelude.rs

@@ -826,6 +826,14 @@ macro_rules! isle_common_prelude_methods {
                 .expect("trusted MemFlagsData not found in DFG")
         }
 
+        #[inline]
+        fn mem_flags_aligned_read_only(&mut self) -> MemFlags {
+            self.dfg()
+                .mem_flags
+                .get(MemFlagsData::new().with_aligned().with_readonly())
+                .expect("aligned, read-only MemFlagsData not found in DFG")
+        }
+
         #[inline]
         fn mem_flags_data(&mut self, flags: MemFlags) -> Option<MemFlagsData> {
             Some(self.dfg().mem_flags[flags])