Add interrupt_poll Cranelift instruction.

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

@@ -370,6 +370,73 @@ fn define_control_flow(
         .call()
         .branches(),
     );
+
+    ig.push(
+        Inst::new(
+            "interrupt_poll",
+            r#"
+        Load from a memory location to possibly trigger a resumeable
+        interruption.
+
+        - Load a pointer-sized value from memory at `load_ptr`, and throw it
+          away.
+        - Keep `context` in a fixed register.
+        - Reserve a second register as scratch space. (Exact choices of
+          registers are ISA-specific; see each backend's `get_operands` for
+          which registers and the reasoning behind them.)
+        - Record the address of the load instruction in the binary's trap table.
+
+        This instruction aids in implementing virtual-memory-triggered
+        interrupts, with the load trapping if the loaded location is
+        inaccessible. The trap handler can then take further action, first using
+        the trap table to confirm that this instruction was the cause. The
+        handler receives further arbitrary input in `context`. It can use the
+        second reserved register as scratch space: for example, to record the
+        original resumption address so it can arrange to "return to" a
+        trampoline first, which can ultimately then jump to the original
+        address. (Such gymnastics are necessary on platforms where signal
+        handlers cannot push stack frames directly.) It is expected that
+        execution will ultimately resume at the load, re-running it; care must
+        be taken to ensure it succeeds the second time, lest the whole process
+        repeat. This is where the output `next_load_ptr` comes in, carrying a
+        new location to load from. All current backends (x64 and aarch64) pin
+        `load_ptr` and `next_load_ptr` to the same register so a move from the
+        latter to the former does not even need to be emitted.
+        "#,
+            &formats.int_add_trap,
+        )
+        .operands_in(&[
+            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"),
+            Operand::new("code", &imm.trapcode)
+                .with_doc("trap code to record at the load's address"),
+        ])
+        .operands_out(&[
+            Operand::new("next_load_ptr", iAddr).with_doc("memory location to load from next time")
+        ])
+        // As with `stack_switch`, this instruction is a call, in that "it
+        // continues execution elsewhere". See reasoning at
+        // https://github.com/bytecodealliance/wasmtime/pull/9078#issuecomment-2273869774.
+        .call()
+        .can_load()
+        // It may transfer control to something that may store. Declaring this
+        // makes us a memory fence.
+        .can_store()
+        // The universe of possible side effects is wide open. Control may never
+        // even return to this point. When this instruction is used to trigger
+        // preemption, we certainly do not want it hoisted or deduplicated via
+        // GVN.
+        .other_side_effects(),
+        // If `load` is not `can_trap()`, this isn't either.
+        //
+        // This cannot use `side_effects_idempotent()`: its purpose is to allow
+        // deduplication or LICM of redundant instructions. The decision made by
+        // `interrupt_poll` (implicitly, via trap) is determined by the flags on
+        // the page of memory loaded, which change at runtime, outside the
+        // purview of the compiler. Thus, the tuple formed by the instruction
+        // and its input operands is not sufficient to decide redundancy.
+    );
 }
 
 #[inline(never)]

cranelift/codegen/src/ir/instructions.rs

@@ -643,9 +643,13 @@ impl InstructionData {
             Self::Ternary {
                 opcode: Opcode::StackSwitch,
                 ..
+            }
+            | Self::IntAddTrap {
+                opcode: Opcode::InterruptPoll,
+                ..
             } => {
-                // `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

@@ -1038,7 +1038,15 @@
        ;; 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, `dst`
+       ;; to x9, and `next_load_ptr` to x10.
+       (InterruptPoll (dst WritableReg)
+                            (load_ptr Reg)
+                            (context Reg)
+                            (trap_code TrapCode)
+                            (next_load_ptr WritableReg))))
 
 (spec (MInst.AluRRImmLogic alu_op size rd rn imml)
       (provide
@@ -5849,6 +5857,7 @@
 (attr f16const (tag TODO))
 (attr fcvt_to_uint_sat (tag TODO))
 (attr get_return_address (tag TODO))
+(attr interrupt_poll (tag TODO))
 (attr lower_bmask (tag TODO))
 (attr nop (tag TODO))
 (attr invalid_reg (tag TODO))

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

@@ -3708,6 +3708,31 @@ impl MachInstEmit for Inst {
                 .emit(sink, emit_info, state);
                 sink.bind_label(loop_end, &mut state.ctrl_plane);
             }
+
+            &Inst::InterruptPoll {
+                dst,
+                load_ptr,
+                trap_code,
+                ..
+            } => {
+                // Record the address of the load in the trap table so a signal
+                // handler can later distinguish whether a segfault is its
+                // fault.
+                sink.add_trap(trap_code);
+
+                // 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);
+            }
         }
 
         let end_off = sink.cur_offset();

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

@@ -968,6 +968,37 @@ fn aarch64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_early_def(start);
             collector.reg_use(end);
         }
+        Inst::InterruptPoll {
+            dst,
+            load_ptr,
+            context,
+            trap_code: _,
+            next_load_ptr,
+        } => {
+            // `load_ptr` is an input param. It is pinned to x10 so an update of
+            // `next_load_ptr` updates this as well. x10 is chosen because it is
+            // a caller-saved scratch reg with no special role.
+            collector.reg_fixed_use(load_ptr, regs::xreg(10));
+            // 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));
+            // `next_load_ptr` is pinned so embedders know where to write to
+            // fill it. It shares x10 with `load_ptr`, above, so filling this
+            // output means also filling in the (potential) next input, for
+            // efficiency. This also means that leaving x10 alone (in the
+            // common, non-trapping case), makes info conceptually flow in the
+            // other direction, piping the old but unchanging `load_ptr` through
+            // to the output.
+            collector.reg_fixed_def(next_load_ptr, regs::xreg(10));
+        }
     }
 }
 
@@ -3014,6 +3045,21 @@ impl Inst {
                 let step = step.pretty_print(0);
                 format!("stack_probe_loop {start}, {end}, {step}")
             }
+            &Inst::InterruptPoll {
+                dst,
+                load_ptr,
+                context,
+                trap_code,
+                next_load_ptr,
+            } => {
+                let dst = pretty_print_reg(dst.to_reg());
+                let load_ptr = pretty_print_reg(load_ptr);
+                let context = pretty_print_reg(context);
+                let next_load_ptr = pretty_print_reg(next_load_ptr.to_reg());
+                format!(
+                    "{next_load_ptr} = interrupt_poll {dst}, {load_ptr}, {context} #trap={trap_code}"
+                )
+            }
         }
     }
 }

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

@@ -2560,6 +2560,20 @@
 (rule (lower (symbol_value _ (symbol_value_data extname dist offset)))
       (load_ext_name (box_external_name extname) offset dist))
 
+;;;; Rules for `interrupt_poll` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(rule (lower (interrupt_poll _ load_ptr context trap_code))
+      (let ((load_ptr Reg (put_in_reg load_ptr))
+            (context Reg (put_in_reg context))
+            (dst WritableReg (temp_writable_reg $I64))
+            (next_load_ptr WritableReg (temp_writable_reg $I64))
+            (_ Unit (emit (MInst.InterruptPoll dst
+                                                     load_ptr
+                                                     context
+                                                     trap_code
+                                                     next_load_ptr))))
+       (output_reg next_load_ptr)))
+
 ;;; Rules for `get_{frame,stack}_pointer` and `get_return_address` ;;;;;;;;;;;;;
 
 (rule (lower (get_frame_pointer _))

cranelift/codegen/src/isa/call_conv.rs

@@ -14,16 +14,21 @@ use serde_derive::{Deserialize, Serialize};
 pub enum CallConv {
     /// Best performance, not ABI-stable.
     Fast,
-    /// Supports tail calls, not ABI-stable except for exception
-    /// payload registers.
+    /// Supports tail calls, not ABI-stable except as stated here.
     ///
-    /// On exception resume, a caller to a `tail`-convention function
-    /// assumes that the exception payload values are in the following
-    /// registers (per platform):
+    /// On exception resume, a caller to a `tail`-convention function assumes
+    /// that the exception payload values are in the following registers (per
+    /// platform):
     /// - x86-64: rax, rdx
     /// - aarch64: x0, x1
     /// - riscv64: a0, a1
     /// - pulley{32,64}: x0, x1
+    ///
+    /// The `interrupt_call` instruction uses registers as follows:
+    /// - x86-64: context in rdi, load ptr and return value in r11, scratch in
+    ///   r10
+    /// - aarch64: context in x0, load ptr and return value in x10, scratch in
+    ///   x9
     //
     // Currently, this is basically sys-v except that callees pop stack
     // arguments, rather than callers. Expected to change even more in the

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,14 @@
                     (offset i64)
                     (distance RelocDistance))
 
+       ;; A load whose result is discarded; `context` is pinned to rdi, `dst`
+       ;; to r10, and `load_ptr`/`next_load_ptr` share r11.
+       (InterruptPoll (dst WritableGpr)
+                            (load_ptr Gpr)
+                            (context Gpr)
+                            (trap_code TrapCode)
+                            (next_load_ptr WritableGpr))
+
        ;; =========================================
        ;; Instructions pertaining to atomic memory accesses.
 

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

@@ -650,6 +650,23 @@ pub(crate) fn emit(
             }
         }
 
+        Inst::InterruptPoll {
+            dst,
+            load_ptr,
+            trap_code,
+            ..
+        } => {
+            // Record the address of load in the trap table so a signal handler
+            // can later distinguish whether a segfault is its fault.
+            sink.add_trap(*trap_code);
+
+            let load_ptr_addr = SyntheticAmode::real(Amode::imm_reg(0, **load_ptr));
+            // Since we're clobbering dst anyway to store the resume 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);
+        }
+
         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::InterruptPoll { .. }
             | Inst::TrapIf { .. }
             | Inst::TrapIfAnd { .. }
             | Inst::TrapIfOr { .. }
@@ -671,6 +672,22 @@ impl PrettyPrint for Inst {
                 )
             }
 
+            Inst::InterruptPoll {
+                dst,
+                load_ptr,
+                context,
+                trap_code,
+                next_load_ptr,
+            } => {
+                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);
+                let next_load_ptr = pretty_print_reg(*next_load_ptr.to_reg(), 8);
+                format!(
+                    "{next_load_ptr} = interrupt_poll {dst}, {load_ptr}, {context} #trap={trap_code}"
+                )
+            }
+
             Inst::JmpKnown { dst } => {
                 let op = ljustify("jmp".to_string());
                 let dst = dst.to_string();
@@ -1051,6 +1068,47 @@ fn x64_get_operands(inst: &mut Inst, collector: &mut impl OperandVisitor) {
             collector.reg_clobbers(clobbers);
         }
 
+        Inst::InterruptPoll {
+            dst,
+            load_ptr,
+            context,
+            trap_code: _,
+            next_load_ptr,
+        } => {
+            // `load_ptr` is an input param. It is pinned to r11 so an update of
+            // `next_load_ptr` updates this as well. r11 is chosen because it is
+            // a caller-saved reg not used for arg-passing in Linux/x64.
+            collector.reg_fixed_use(load_ptr, regs::r11());
+            // Demand context (vmctx) go into RDI.
+            collector.reg_fixed_use(context, regs::rdi());
+            // Reserve r10 as a place for a signal handler to stow the original
+            // resume address. This allows the handler to twiddle saved machine
+            // state to return to a custom trampoline when it exits, allowing it
+            // to accomplish things that are unsafe at interrupt time. The
+            // trampoline can jump to the address in r10 when done to resume.
+            //
+            // r10 is chosen because it is 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 Cranelift
+            // does not emit those. (Do take care if you are interacting with
+            // external ones.) It is also used to store the function stack limit
+            // in Cranelift, but the stack-limit check is confined to the
+            // function prologue and thus over by the time this instruction is
+            // used.
+            //
+            // 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());
+            // `next_load_ptr` is pinned so embedders know where to write to
+            // fill it. It shares r11 with `load_ptr`, above, so filling this
+            // output means also filling in the (potential) next input, for
+            // efficiency. This also means that leaving r11 alone (in the
+            // common, non-trapping case), makes info conceptually flow in the
+            // other direction, piping the old but unchanging `load_ptr` through
+            // to the output.
+            collector.reg_fixed_def(next_load_ptr, regs::r11());
+        }
+
         Inst::ReturnCallKnown { info } => {
             let ReturnCallInfo {
                 dest, uses, tmp, ..