Replace nops with dead loads from the interrupt page. Add MMU-based epoch checks to loop headers as well.

Cache the interrupt page ptr in a local for speed, as we did with the epoch deadline.

Here is how I interpret the generated code in epoch-interruption-mmu.wat:
```
;; @001B                               v2 = load.i64 notrap aligned readonly can_move v0+8 ;; Skip over magic number (4b) and alignment (another 4b).
;; @001B                               v3 = load.i64 notrap aligned v2+16 ;; Get interrupt page ptr.
;; @001B                               v4 = load.i32 aligned readonly v3 ;; Read from page ptr.
```

Author: erikrose

crates/cranelift/src/func_environ.rs

@@ -153,6 +153,10 @@ pub struct FuncEnvironment<'module_environment> {
     /// any yield.
     epoch_deadline_var: cranelift_frontend::Variable,
 
+    // A cached pointer to the epoch interrupt page so we don't have to
+    // continually dig it out of the `VMStoreContext`
+    epoch_interrupt_page_ptr_var: cranelift_frontend::Variable,
+
     /// A cached pointer to the per-Engine epoch counter, when
     /// performing epoch-based interruption. Initialized in the
     /// function prologue. We prefer to use a variable here rather
@@ -216,6 +220,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             fuel_var: Variable::reserved_value(),
             epoch_deadline_var: Variable::reserved_value(),
             epoch_ptr_var: Variable::reserved_value(),
+            epoch_interrupt_page_ptr_var: Variable::reserved_value(),
 
             // Start with at least one fuel being consumed because even empty
             // functions should consume at least some fuel.
@@ -589,6 +594,43 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         self.epoch_check_full(builder, cur_epoch_value, continuation_block);
     }
 
+    /// Codegens what needs to go at the top of a function to support
+    /// epoch_interrupt_via_mmu.
+    fn epoch_mmu_function_entry(&mut self, builder: &mut FunctionBuilder<'_>) {
+        debug_assert!(self.epoch_interrupt_page_ptr_var.is_reserved_value());
+        self.epoch_interrupt_page_ptr_var = builder.declare_var(self.pointer_type());
+
+        // Cache ptr to interrupt page in a local (and hopefully a register, at
+        // the discretion of regalloc), rather than digging it out of the
+        // `VMStoreContext` every time.
+        let vmstore_ctx = self.get_vmstore_context_ptr(builder);
+        let epoch_interrupt_page_ptr = builder.ins().load(
+            self.pointer_type(),
+            ir::MemFlags::trusted(),
+            vmstore_ctx,
+            ir::immediates::Offset32::new(
+                self.offsets.ptr.vmstore_context_epoch_interrupt_page_ptr() as i32,
+            ),
+        );
+        builder.def_var(self.epoch_interrupt_page_ptr_var, epoch_interrupt_page_ptr);
+
+        Self::epoch_mmu_interruption_check(epoch_interrupt_page_ptr, builder);
+    }
+
+    /// Codegens a dead load from the epoch interrupt page, which causes a trap
+    /// if an interrupt is due.
+    fn epoch_mmu_interruption_check(
+        epoch_interrupt_page_ptr: ir::Value,
+        builder: &mut FunctionBuilder<'_>,
+    ) {
+        let _ = builder.ins().load(
+            ir::types::I32, // Arbitrary. Pick whatever works on all ISAs and is fastest.
+            ir::MemFlags::new().with_aligned().with_readonly(),
+            epoch_interrupt_page_ptr,
+            ir::immediates::Offset32::new(0),
+        );
+    }
+
     #[cfg(feature = "wmemcheck")]
     fn hook_malloc_exit(&mut self, builder: &mut FunctionBuilder, retvals: &[ir::Value]) {
         let check_malloc = self.builtin_functions.check_malloc(builder.func);
@@ -3314,6 +3356,13 @@ impl FuncEnvironment<'_> {
             self.epoch_check(builder);
         }
 
+        // If we're using MMU-based epoch detection, provoke an interrupt if
+        // it's time.
+        if self.tunables.epoch_interruption_via_mmu {
+            let page_ptr = builder.use_var(self.epoch_interrupt_page_ptr_var);
+            Self::epoch_mmu_interruption_check(page_ptr, builder);
+        }
+
         Ok(())
     }
 
@@ -3375,8 +3424,7 @@ impl FuncEnvironment<'_> {
         }
 
         if self.tunables.epoch_interruption_via_mmu {
-            builder.ins().iconst(I32, 33); // a useless constant, hopefully not optimized out
-            // NEXT: Dead-load something from the vmctx instead.
+            self.epoch_mmu_function_entry(builder);
         }
 
         #[cfg(feature = "wmemcheck")]

crates/environ/src/vmoffsets.rs

@@ -175,10 +175,16 @@ pub trait PtrSize {
         self.vmstore_context_fuel_consumed() + 8
     }
 
+    /// Return the offset of the `epoch_interrupt_page_ptr` field of `VMStoreContext`.
+    #[inline]
+    fn vmstore_context_epoch_interrupt_page_ptr(&self) -> u8 {
+        self.vmstore_context_epoch_deadline() + 8
+    }
+
     /// Return the offset of the `stack_limit` field of `VMStoreContext`
     #[inline]
     fn vmstore_context_stack_limit(&self) -> u8 {
-        self.vmstore_context_epoch_deadline() + 8
+        self.vmstore_context_epoch_interrupt_page_ptr() + self.size()
     }
 
     /// Return the offset of the `gc_heap` field of `VMStoreContext`.

tests/disas/epoch-interruption-mmu.wat

@@ -8,11 +8,15 @@
 ;; function u0:0(i64 vmctx, i64) tail {
 ;;     gv0 = vmctx
 ;;     gv1 = load.i64 notrap aligned readonly gv0+8
-;;     gv2 = load.i64 notrap aligned gv1+16
+;;     gv2 = load.i64 notrap aligned gv1+24
+;;     gv3 = vmctx
+;;     gv4 = load.i64 notrap aligned readonly can_move gv3+8
 ;;     stack_limit = gv2
 ;;
 ;;                                 block0(v0: i64, v1: i64):
-;; @001b                               v2 = iconst.i32 33
+;; @001b                               v2 = load.i64 notrap aligned readonly can_move v0+8
+;; @001b                               v3 = load.i64 notrap aligned v2+16
+;; @001b                               v4 = load.i32 aligned readonly v3
 ;; @001c                               jump block1
 ;;
 ;;                                 block1: