Switch interpreter dispatch to musttail (become) threading

.github/workflows/rust.yml

@@ -143,3 +143,12 @@ 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
+    continue-on-error: true
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install latest nightly
+      run: rustup toolchain install nightly --profile minimal
+    - name: Build and test with interpreter-musttail-dispatch
+      run: ./ci/jobs/build-and-test-nightly-musttail.sh

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

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+set -euo pipefail
+cd -- "$(dirname -- "${BASH_SOURCE[0]}")"
+cd ../..
+
+echo ">> cargo +nightly test (polkavm, interpreter-musttail-dispatch, interpreter tests only)"
+cargo +nightly test -p polkavm --features interpreter-musttail-dispatch -- tests::interpreter_
+
+echo ">> cargo +nightly test --release (polkavm, interpreter-musttail-dispatch, interpreter tests only)"
+cargo +nightly test --release -p polkavm --features interpreter-musttail-dispatch -- tests::interpreter_

crates/polkavm/Cargo.toml

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

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 = "interpreter-musttail-dispatch")]
+type HandlerResult = InterruptKind;
+#[cfg(not(feature = "interpreter-musttail-dispatch"))]
 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 = "interpreter-musttail-dispatch")]
+        {
+            dispatch::<DEBUG>(self, self.next_compiled_offset)
+        }
+        #[cfg(not(feature = "interpreter-musttail-dispatch"))]
+        {
+            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,44 @@ struct Args {
 
 type Handler = for<'a> fn(visitor: &'a mut InterpretedInstance, compiled_offset: Target) -> HandlerResult;
 
+// `become` is parse-gated; hiding it in a macro keeps the parser happy when the feature is off.
+#[cfg(feature = "interpreter-musttail-dispatch")]
+macro_rules! tail_call {
+    ($e:expr) => {
+        become $e
+    };
+}
+
+#[cfg(feature = "interpreter-musttail-dispatch")]
+#[inline(always)]
+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()) {
+        tail_call!(handler(visitor, off))
+    } else {
+        visitor.interrupt.clone()
+    }
+}
+
+// Tail-jump from a handler into the next one. `become` lives at the handler tail so the chain
+// stays musttail even if `dispatch` isn't inlined (e.g. at `-O0`). With the feature off, the
+// handler returns the new offset and the central loop drives the next call.
+#[cfg(feature = "interpreter-musttail-dispatch")]
+macro_rules! tail_dispatch {
+    ($self:expr, $next_off:expr) => {
+        tail_call!(dispatch::<DEBUG>($self, $next_off))
+    };
+}
+
+#[cfg(not(feature = "interpreter-musttail-dispatch"))]
+macro_rules! tail_dispatch {
+    ($self:expr, $next_off:expr) => {
+        $next_off
+    };
+}
+
 macro_rules! define_interpreter {
     (@define $handler_name:ident $body:block $self:ident $compiled_offset:ident) => {{
         impl Args {
@@ -2949,7 +2998,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/lib.rs

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