Switch interpreter dispatch to musttail (become) threading

.github/workflows/rust.yml

@@ -143,3 +143,11 @@ jobs:
       run: rustup run nightly-2024-07-10 cargo install --locked --version 0.12.0 cargo-fuzz
     - name: Fuzz
       run: ./ci/jobs/fuzz.sh
+  nightly-musttail:
+    runs-on: ubuntu-22.04
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install latest nightly
+      run: rustup toolchain install nightly --profile minimal
+    - name: Build and test with experimental-musttail
+      run: ./ci/jobs/build-and-test-nightly-musttail.sh

ci/jobs/build-and-test-nightly-musttail.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+set -euo pipefail
+cd -- "$(dirname -- "${BASH_SOURCE[0]}")"
+cd ../..
+
+if [ "$(cat /proc/sys/vm/unprivileged_userfaultfd)" != "1" ]; then
+    echo "1" | sudo tee /proc/sys/vm/unprivileged_userfaultfd > /dev/null
+fi
+
+echo ">> cargo +nightly build (polkavm, experimental-musttail)"
+cargo +nightly build -p polkavm --features experimental-musttail
+
+echo ">> cargo +nightly test (polkavm, experimental-musttail)"
+cargo +nightly test -p polkavm --features experimental-musttail
+
+echo ">> cargo +nightly test --release (polkavm, experimental-musttail)"
+cargo +nightly test --release -p polkavm --features experimental-musttail

crates/polkavm/Cargo.toml

@@ -50,3 +50,6 @@ generic-sandbox = []
 
 # Internal feature for testing. DO NOT USE.
 export-internals-for-testing = []
+
+# Musttail (`become`) interpreter dispatch. Requires nightly.
+experimental-musttail = []

crates/polkavm/src/interpreter.rs

@@ -25,6 +25,10 @@ use polkavm_common::program::{
 use polkavm_common::utils::{align_to_next_page_usize, slice_assume_init_mut, ArcBytes, GasVisitorT};
 
 type Target = u32;
+
+#[cfg(feature = "experimental-musttail")]
+type HandlerResult = InterruptKind;
+#[cfg(not(feature = "experimental-musttail"))]
 type HandlerResult = Target;
 
 #[derive(Copy, Clone)]
@@ -1711,16 +1715,23 @@ impl InterpretedInstance {
             log::trace!("Implicitly resuming at: [{}]", self.next_compiled_offset);
         }
 
-        let mut offset = self.next_compiled_offset;
-        loop {
-            if DEBUG {
-                self.cycle_counter += 1;
-            }
+        #[cfg(feature = "experimental-musttail")]
+        {
+            dispatch::<DEBUG>(self, self.next_compiled_offset)
+        }
+        #[cfg(not(feature = "experimental-musttail"))]
+        {
+            let mut offset = self.next_compiled_offset;
+            loop {
+                if DEBUG {
+                    self.cycle_counter += 1;
+                }
 
-            if let Some(handler) = self.compiled_handlers.get(cast(offset).to_usize()) {
-                offset = handler(self, offset);
-            } else {
-                return self.interrupt.clone();
+                if let Some(handler) = self.compiled_handlers.get(cast(offset).to_usize()) {
+                    offset = handler(self, offset);
+                } else {
+                    return self.interrupt.clone();
+                }
             }
         }
     }
@@ -2504,6 +2515,36 @@ struct Args {
 
 type Handler = for<'a> fn(visitor: &'a mut InterpretedInstance, compiled_offset: Target) -> HandlerResult;
 
+// `become` is parse-time feature-gated, so the musttail dispatcher must live in a separately-loaded file.
+#[cfg(feature = "experimental-musttail")]
+#[path = "interpreter_musttail.rs"]
+mod musttail;
+
+#[cfg(feature = "experimental-musttail")]
+use musttail::{dispatch, handler_tail};
+
+#[cfg(not(feature = "experimental-musttail"))]
+#[inline(always)]
+fn handler_tail<const DEBUG: bool>(_visitor: &mut InterpretedInstance, next_off: Target) -> HandlerResult {
+    next_off
+}
+
+// `become` from the handler itself, so the tail call doesn't depend on
+// `handler_tail`/`dispatch` being inlined (which `-O0` may skip).
+#[cfg(feature = "experimental-musttail")]
+macro_rules! tail_dispatch {
+    ($self:expr, $next_off:expr) => {
+        become handler_tail::<DEBUG>($self, $next_off)
+    };
+}
+
+#[cfg(not(feature = "experimental-musttail"))]
+macro_rules! tail_dispatch {
+    ($self:expr, $next_off:expr) => {
+        handler_tail::<DEBUG>($self, $next_off)
+    };
+}
+
 macro_rules! define_interpreter {
     (@define $handler_name:ident $body:block $self:ident $compiled_offset:ident) => {{
         impl Args {
@@ -2949,7 +2990,8 @@ macro_rules! define_interpreter {
                 #[allow(clippy::needless_lifetimes)]
                 pub fn $handler_name<'a, $(M: $M_ty,)? $(const $const: $const_ty),+>($self: &'a mut InterpretedInstance, compiled_offset: Target) -> HandlerResult {
                     let $compiled_offset = compiled_offset;
-                    define_interpreter!(@define $handler_name $body $self $compiled_offset $($arg)*)
+                    let next_off: Target = define_interpreter!(@define $handler_name $body $self $compiled_offset $($arg)*);
+                    tail_dispatch!($self, next_off)
                 }
             )+
         }

crates/polkavm/src/interpreter_musttail.rs

@@ -0,0 +1,19 @@
+use super::{HandlerResult, InterpretedInstance, Target};
+use polkavm_common::cast::cast;
+
+#[inline(always)]
+pub(super) fn handler_tail<const DEBUG: bool>(visitor: &mut InterpretedInstance, next_off: Target) -> HandlerResult {
+    become dispatch::<DEBUG>(visitor, next_off);
+}
+
+#[inline(always)]
+pub(super) fn dispatch<const DEBUG: bool>(visitor: &mut InterpretedInstance, off: Target) -> HandlerResult {
+    if DEBUG {
+        visitor.cycle_counter += 1;
+    }
+    if let Some(&handler) = visitor.compiled_handlers.get(cast(off).to_usize()) {
+        become handler(visitor, off);
+    } else {
+        visitor.interrupt.clone()
+    }
+}

crates/polkavm/src/lib.rs

@@ -1,4 +1,6 @@
 #![cfg_attr(not(feature = "std"), no_std)]
+#![cfg_attr(feature = "experimental-musttail", feature(explicit_tail_calls))]
+#![cfg_attr(feature = "experimental-musttail", allow(incomplete_features))]
 #![forbid(unused_must_use)]
 #![forbid(clippy::missing_safety_doc)]
 #![deny(clippy::undocumented_unsafe_blocks)]