feat(jit): add tail side-entry link ABI

fffonion · fffonion · commit 4cf464a95074 · 2026-07-19T01:25:55.000+08:00
diff --git a/src/vm/jit/native/lower.rs b/src/vm/jit/native/lower.rs
@@ -37,9 +37,10 @@ use crate::vm::native::{
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::immediates::Ieee64;
 use cranelift_codegen::ir::{
-    Block, BlockArg, InstBuilder, MemFlags, StackSlot, StackSlotData, StackSlotKind, types,
+    AbiParam, Block, BlockArg, InstBuilder, MemFlags, Signature, StackSlot, StackSlotData,
+    StackSlotKind, types,
 };
-use cranelift_codegen::isa::OwnedTargetIsa;
+use cranelift_codegen::isa::{CallConv, OwnedTargetIsa};
 use cranelift_codegen::settings::{self, Configurable};
 use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
 use cranelift_jit::{JITBuilder, JITModule};
@@ -50,6 +51,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
 
 static CRANELIFT_TRACE_ID: AtomicU64 = AtomicU64::new(1);
 static CRANELIFT_JIT_ISA: OnceLock<Result<OwnedTargetIsa, String>> = OnceLock::new();
+static CRANELIFT_TAIL_ISA: OnceLock<Result<OwnedTargetIsa, String>> = OnceLock::new();
 
 type TaggedConstants = (Box<[Value]>, HashMap<SsaValueId, usize>);
 
@@ -78,6 +80,148 @@ impl TraceKeepAlive {
     }
 }
 
+pub(crate) struct CompiledTailFunction {
+    entry: *const u8,
+    _keepalive: TraceKeepAlive,
+    code: Vec<u8>,
+}
+
+impl CompiledTailFunction {
+    pub(crate) fn entry(&self) -> *const u8 {
+        self.entry
+    }
+
+    pub(crate) fn code_len(&self) -> usize {
+        self.code.len()
+    }
+}
+
+fn tail_entry_signature(pointer_type: cranelift_codegen::ir::Type) -> Signature {
+    let mut signature = Signature::new(CallConv::Tail);
+    signature.params.push(AbiParam::new(pointer_type));
+    signature.returns.push(AbiParam::new(types::I32));
+    signature
+}
+
+fn compile_standalone_native_function(
+    prefix: &str,
+    signature: impl FnOnce(cranelift_codegen::ir::Type, CallConv) -> Signature,
+    lower: impl FnOnce(&mut FunctionBuilder<'_>, cranelift_codegen::ir::Type, CallConv) -> VmResult<()>,
+) -> VmResult<CompiledTailFunction> {
+    let isa = native_tail_isa()?;
+    let jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
+    let mut module = JITModule::new(jit_builder);
+    let pointer_type = module.target_config().pointer_type();
+    let default_call_conv = module.target_config().default_call_conv;
+    let mut ctx = module.make_context();
+    ctx.func.signature = signature(pointer_type, default_call_conv);
+    let function_id = CRANELIFT_TRACE_ID.fetch_add(1, Ordering::Relaxed);
+    let function_name = format!("{prefix}_{function_id}");
+    let func_id = module
+        .declare_function(&function_name, Linkage::Local, &ctx.func.signature)
+        .map_err(|err| VmError::JitNative(format!("declare {prefix} failed: {err}")))?;
+    {
+        let mut fb_ctx = FunctionBuilderContext::new();
+        let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fb_ctx);
+        lower(&mut builder, pointer_type, default_call_conv)?;
+        builder.seal_all_blocks();
+        builder.finalize();
+    }
+    module
+        .define_function(func_id, &mut ctx)
+        .map_err(|err| VmError::JitNative(format!("define {prefix} failed: {err}")))?;
+    let code_len = ctx
+        .compiled_code()
+        .ok_or_else(|| VmError::JitNative(format!("{prefix} produced no machine code")))?
+        .code_buffer()
+        .len();
+    module.clear_context(&mut ctx);
+    module
+        .finalize_definitions()
+        .map_err(|err| VmError::JitNative(format!("finalize {prefix} failed: {err}")))?;
+    let module_entry = module.get_finalized_function(func_id);
+    let code = unsafe { std::slice::from_raw_parts(module_entry, code_len).to_vec() };
+    let keepalive = TraceKeepAlive::from_code(&code, Box::new([]))?;
+    let entry = keepalive.entry();
+    Ok(CompiledTailFunction {
+        entry,
+        _keepalive: keepalive,
+        code,
+    })
+}
+
+pub(crate) fn compile_tail_status_body(status: i32) -> VmResult<CompiledTailFunction> {
+    compile_standalone_native_function(
+        "pd_vm_tail_status",
+        |pointer_type, _| tail_entry_signature(pointer_type),
+        move |builder, _, _| {
+            let entry = builder.create_block();
+            builder.append_block_params_for_function_params(entry);
+            builder.switch_to_block(entry);
+            let status = builder.ins().iconst(types::I32, i64::from(status));
+            builder.ins().return_(&[status]);
+            Ok(())
+        },
+    )
+}
+
+pub(crate) fn compile_tail_side_link_body(
+    slot_address: usize,
+    deopt_status: i32,
+) -> VmResult<CompiledTailFunction> {
+    compile_standalone_native_function(
+        "pd_vm_tail_side_link",
+        |pointer_type, _| tail_entry_signature(pointer_type),
+        move |builder, pointer_type, _| {
+            let entry = builder.create_block();
+            let deopt = builder.create_block();
+            let linked = builder.create_block();
+            builder.append_block_params_for_function_params(entry);
+            builder.switch_to_block(entry);
+            let vm_ptr = builder.block_params(entry)[0];
+            let slot_address = iconst_ptr_from_addr(builder, pointer_type, slot_address)?;
+            let target = builder
+                .ins()
+                .load(pointer_type, MemFlags::new(), slot_address, 0);
+            let is_null = builder.ins().icmp_imm(IntCC::Equal, target, 0);
+            builder.ins().brif(is_null, deopt, &[], linked, &[]);
+
+            builder.switch_to_block(deopt);
+            let status = builder.ins().iconst(types::I32, i64::from(deopt_status));
+            builder.ins().return_(&[status]);
+
+            builder.switch_to_block(linked);
+            let signature = builder.import_signature(tail_entry_signature(pointer_type));
+            builder
+                .ins()
+                .return_call_indirect(signature, target, &[vm_ptr]);
+            Ok(())
+        },
+    )
+}
+
+pub(crate) fn compile_system_tail_wrapper(root_entry: *const u8) -> VmResult<CompiledTailFunction> {
+    let root_entry = root_entry as usize;
+    compile_standalone_native_function(
+        "pd_vm_tail_wrapper",
+        entry_signature,
+        move |builder, pointer_type, _| {
+            let entry = builder.create_block();
+            builder.append_block_params_for_function_params(entry);
+            builder.switch_to_block(entry);
+            let vm_ptr = builder.block_params(entry)[0];
+            let root_entry = iconst_ptr_from_addr(builder, pointer_type, root_entry)?;
+            let signature = builder.import_signature(tail_entry_signature(pointer_type));
+            let call = builder
+                .ins()
+                .call_indirect(signature, root_entry, &[vm_ptr]);
+            let status = builder.inst_results(call)[0];
+            builder.ins().return_(&[status]);
+            Ok(())
+        },
+    )
+}
+
 fn try_compile_ssa_trace(
     trace: &JitTrace,
     ssa: &SsaTrace,
@@ -4950,6 +5094,24 @@ pub(crate) fn compile_trace(
     })
 }
 
+fn native_tail_isa() -> VmResult<OwnedTargetIsa> {
+    let cached = CRANELIFT_TAIL_ISA.get_or_init(|| {
+        let mut flag_builder = settings::builder();
+        flag_builder
+            .set("opt_level", "speed")
+            .map_err(|err| format!("failed to set cranelift opt_level: {err}"))?;
+        flag_builder
+            .set("preserve_frame_pointers", "true")
+            .map_err(|err| format!("failed to preserve tail-call frame pointers: {err}"))?;
+        let isa_builder = cranelift_native::builder()
+            .map_err(|err| format!("failed to build native tail ISA: {err}"))?;
+        isa_builder
+            .finish(settings::Flags::new(flag_builder))
+            .map_err(|err| format!("failed to finalize cranelift tail ISA: {err}"))
+    });
+    cached.clone().map_err(VmError::JitNative)
+}
+
 fn native_isa(profile: NativeCompileProfile) -> VmResult<OwnedTargetIsa> {
     let cached = match profile {
         NativeCompileProfile::Jit => &CRANELIFT_JIT_ISA,
diff --git a/src/vm/jit/native/mod.rs b/src/vm/jit/native/mod.rs
@@ -2,6 +2,9 @@
 #[cfg(not(feature = "cranelift-jit"))]
 use super::super::super::VmError;
 use super::super::super::VmResult;
+use super::ir::{SsaMaterialization, SsaValueRepr};
+use crate::ValueType;
+use std::sync::atomic::{AtomicPtr, Ordering};
 
 pub(crate) use crate::vm::native::{
     NativeInterruptSettings, STATUS_CONTINUE, STATUS_ERROR, STATUS_HALTED, STATUS_LINKED_CONTINUE,
@@ -30,6 +33,88 @@ impl TraceLoweringKind {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum SideEntryOwnership {
+    Borrowed,
+    Owned,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum InheritedStateAbiClass {
+    ScalarInt,
+    ScalarFloat,
+    ScalarBool,
+    HeapPointer {
+        tag: ValueType,
+        ownership: SideEntryOwnership,
+    },
+    Tagged(SideEntryOwnership),
+}
+
+pub(crate) fn classify_side_entry_repr(
+    repr: SsaValueRepr,
+    ownership: SideEntryOwnership,
+) -> InheritedStateAbiClass {
+    match repr {
+        SsaValueRepr::I64 => InheritedStateAbiClass::ScalarInt,
+        SsaValueRepr::F64 => InheritedStateAbiClass::ScalarFloat,
+        SsaValueRepr::Bool => InheritedStateAbiClass::ScalarBool,
+        SsaValueRepr::HeapPtr(tag) => InheritedStateAbiClass::HeapPointer { tag, ownership },
+        SsaValueRepr::Tagged => InheritedStateAbiClass::Tagged(ownership),
+    }
+}
+
+pub(crate) fn classify_side_entry_materialization(
+    materialization: &SsaMaterialization,
+) -> InheritedStateAbiClass {
+    match materialization {
+        SsaMaterialization::Value(_) => {
+            InheritedStateAbiClass::Tagged(SideEntryOwnership::Borrowed)
+        }
+        SsaMaterialization::BoxInt(_) => InheritedStateAbiClass::ScalarInt,
+        SsaMaterialization::BoxFloat(_) => InheritedStateAbiClass::ScalarFloat,
+        SsaMaterialization::BoxBool(_) => InheritedStateAbiClass::ScalarBool,
+        SsaMaterialization::BoxHeapPtr { tag, .. } => InheritedStateAbiClass::HeapPointer {
+            tag: *tag,
+            ownership: SideEntryOwnership::Owned,
+        },
+    }
+}
+
+pub(crate) struct NativeSideLinkSlot {
+    entry: AtomicPtr<u8>,
+}
+
+impl NativeSideLinkSlot {
+    pub(crate) const fn new() -> Self {
+        Self {
+            entry: AtomicPtr::new(std::ptr::null_mut()),
+        }
+    }
+
+    pub(crate) fn target(&self) -> *mut u8 {
+        self.entry.load(Ordering::Acquire)
+    }
+
+    pub(crate) fn publish(&self, entry: *const u8) {
+        self.entry.store(entry.cast_mut(), Ordering::Release);
+    }
+
+    pub(crate) fn clear(&self) {
+        self.entry.store(std::ptr::null_mut(), Ordering::Release);
+    }
+
+    pub(crate) fn address(&self) -> *mut *mut u8 {
+        self.entry.as_ptr()
+    }
+}
+
+impl Default for NativeSideLinkSlot {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 #[cfg(feature = "cranelift-jit")]
 pub(crate) use lower::{CompiledTrace, TraceKeepAlive};
 
@@ -102,7 +187,89 @@ pub(super) fn compile_native_region(
 
 #[cfg(test)]
 mod tests {
-    use super::selected_codegen_backend;
+    use super::lower::{
+        compile_system_tail_wrapper, compile_tail_side_link_body, compile_tail_status_body,
+    };
+    use super::{
+        InheritedStateAbiClass, NativeSideLinkSlot, SideEntryOwnership,
+        classify_side_entry_materialization, classify_side_entry_repr, selected_codegen_backend,
+    };
+    use crate::ValueType;
+    use crate::vm::jit::ir::{SsaMaterialization, SsaValueId, SsaValueRepr};
+
+    #[test]
+    fn side_entry_abi_classifies_scalar_pointer_tagged_borrowed_and_owned_values() {
+        let value = SsaValueId::new(7);
+        assert_eq!(
+            classify_side_entry_materialization(&SsaMaterialization::BoxInt(value)),
+            InheritedStateAbiClass::ScalarInt
+        );
+        assert_eq!(
+            classify_side_entry_materialization(&SsaMaterialization::BoxFloat(value)),
+            InheritedStateAbiClass::ScalarFloat
+        );
+        assert_eq!(
+            classify_side_entry_materialization(&SsaMaterialization::BoxBool(value)),
+            InheritedStateAbiClass::ScalarBool
+        );
+        assert_eq!(
+            classify_side_entry_materialization(&SsaMaterialization::Value(value)),
+            InheritedStateAbiClass::Tagged(SideEntryOwnership::Borrowed)
+        );
+        assert_eq!(
+            classify_side_entry_materialization(&SsaMaterialization::BoxHeapPtr {
+                value,
+                tag: ValueType::String,
+            }),
+            InheritedStateAbiClass::HeapPointer {
+                tag: ValueType::String,
+                ownership: SideEntryOwnership::Owned,
+            }
+        );
+        assert_eq!(
+            classify_side_entry_repr(
+                SsaValueRepr::HeapPtr(ValueType::Array),
+                SideEntryOwnership::Borrowed,
+            ),
+            InheritedStateAbiClass::HeapPointer {
+                tag: ValueType::Array,
+                ownership: SideEntryOwnership::Borrowed,
+            }
+        );
+    }
+
+    #[cfg(feature = "cranelift-jit")]
+    #[test]
+    fn trace_jit_side_link_slot_switches_between_deopt_and_child() {
+        if selected_codegen_backend() != "native" {
+            return;
+        }
+        const DEOPT_STATUS: i32 = 17;
+        const CHILD_STATUS: i32 = 23;
+        let slot = Box::new(NativeSideLinkSlot::new());
+        let child = compile_tail_status_body(CHILD_STATUS).expect("tail child should compile");
+        let root = compile_tail_side_link_body(slot.address() as usize, DEOPT_STATUS)
+            .expect("tail root should compile");
+        let wrapper =
+            compile_system_tail_wrapper(root.entry()).expect("system wrapper should compile");
+        assert!(child.code_len() > 0);
+        assert!(root.code_len() > 0);
+        assert!(wrapper.code_len() > 0);
+        let entry = unsafe {
+            std::mem::transmute::<*const u8, unsafe extern "C" fn(*mut crate::Vm) -> i32>(
+                wrapper.entry(),
+            )
+        };
+
+        assert!(slot.target().is_null());
+        assert_eq!(unsafe { entry(std::ptr::null_mut()) }, DEOPT_STATUS);
+        slot.publish(child.entry());
+        assert_eq!(slot.target().cast_const(), child.entry());
+        assert_eq!(unsafe { entry(std::ptr::null_mut()) }, CHILD_STATUS);
+        slot.clear();
+        assert!(slot.target().is_null());
+        assert_eq!(unsafe { entry(std::ptr::null_mut()) }, DEOPT_STATUS);
+    }
 
     #[test]
     fn selected_backend_is_native() {
