pd-vm: codex: lower ldloc.copy directly in jit traces

fffonion · fffonion · commit 3c0cb0099c28 · 2026-03-10T06:37:54.000+08:00
diff --git a/pd-vm/examples/mini_bench.rs b/pd-vm/examples/mini_bench.rs
@@ -8,9 +8,9 @@ use std::time::{Duration, Instant};
 
 use serde_json::json;
 use vm::{
-    CallOutcome, CompiledProgram, HostFunction, HostFunctionRegistry, JitConfig, Program,
-    SourceFlavor, Value, Vm, VmError, VmStatus, compile_source, compile_source_file,
-    compile_source_with_flavor,
+    CallOutcome, CompiledProgram, HostFunction, HostFunctionRegistry, JitConfig, JitSnapshot,
+    JitTraceTerminal, Program, SourceFlavor, Value, Vm, VmError, VmStatus, compile_source,
+    compile_source_file, compile_source_with_flavor,
 };
 
 const DEFAULT_COMPILE_ITERS: usize = 20;
@@ -767,7 +767,6 @@ fn measure_pd_vm_lua_compare_mode(
         hot_loop_threshold: 1,
         max_trace_len: 16_384,
     });
-    vm.set_jit_runtime_diagnostics_enabled(jit_enabled);
     vm.set_jit_native_bridge_stats_enabled(jit_enabled);
     vm.set_interpreter_opcode_profiling_enabled(!jit_enabled);
 
@@ -777,12 +776,12 @@ fn measure_pd_vm_lua_compare_mode(
         .map_err(|err| format!("warmup {} run failed: {err}", mode.label()))?;
     ensure_expected_completion(&vm, warm_status, &expected_stack, mode.label())?;
     vm.reset_for_reuse();
-    vm.clear_jit_runtime_diagnostics();
     vm.clear_jit_native_bridge_stats();
     vm.clear_interpreter_opcode_profile();
 
     let native_trace_count_before = vm.jit_native_trace_count();
     let native_exec_before = vm.jit_native_exec_count();
+    let jit_snapshot_before = vm.jit_snapshot();
     let started = Instant::now();
     let total_program_runs = timed_runs.saturating_mul(outer as usize);
     for run_index in 0..total_program_runs {
@@ -798,6 +797,7 @@ fn measure_pd_vm_lua_compare_mode(
         }
     }
     let total_elapsed = started.elapsed();
+    let jit_snapshot_after = vm.jit_snapshot();
 
     Ok(LuaCompareSample {
         mode,
@@ -817,7 +817,10 @@ fn measure_pd_vm_lua_compare_mode(
         native_trace_count_before: Some(native_trace_count_before),
         native_trace_count_after: Some(vm.jit_native_trace_count()),
         native_exec_delta: Some(vm.jit_native_exec_count().saturating_sub(native_exec_before)),
-        jit_runtime_diagnostics: vm.jit_runtime_diagnostics_snapshot(),
+        jit_runtime_diagnostics: trace_execution_diagnostics(
+            &jit_snapshot_before,
+            &jit_snapshot_after,
+        ),
         native_bridge_stats: vm
             .jit_native_bridge_stats_snapshot()
             .into_iter()
@@ -997,6 +1000,40 @@ fn normalized_ns_per_inner_iter(
     elapsed.as_nanos() as f64 / (inner as f64 * outer as f64 * timed_runs as f64)
 }
 
+fn trace_execution_diagnostics(
+    before: &JitSnapshot,
+    after: &JitSnapshot,
+) -> Vec<(String, u64)> {
+    let mut before_execs = before
+        .traces
+        .iter()
+        .map(|trace| (trace.id, trace.executions))
+        .collect::<std::collections::HashMap<_, _>>();
+    let mut loop_back = 0_u64;
+    let mut branch_exit = 0_u64;
+    let mut halt = 0_u64;
+    for trace in &after.traces {
+        let before_exec = before_execs.remove(&trace.id).unwrap_or(0);
+        let delta = trace.executions.saturating_sub(before_exec);
+        match trace.terminal {
+            JitTraceTerminal::LoopBack => loop_back = loop_back.saturating_add(delta),
+            JitTraceTerminal::BranchExit => branch_exit = branch_exit.saturating_add(delta),
+            JitTraceTerminal::Halt => halt = halt.saturating_add(delta),
+        }
+    }
+    let mut entries = Vec::new();
+    for (name, count) in [
+        ("loop_back_trace_exec_delta", loop_back),
+        ("branch_exit_trace_exec_delta", branch_exit),
+        ("halt_trace_exec_delta", halt),
+    ] {
+        if count > 0 {
+            entries.push((name.to_string(), count));
+        }
+    }
+    entries
+}
+
 fn write_lua_compare_artifacts(
     config: &BenchConfig,
     shared_source: &str,
diff --git a/pd-vm/src/vm/jit/aot.rs b/pd-vm/src/vm/jit/aot.rs
@@ -602,6 +602,10 @@ fn encode_trace_step(step: &TraceStep, out: &mut Vec<u8>) -> VmResult<()> {
             out.push(17);
             out.push(*index);
         }
+        TraceStep::LdlocCopy(index) => {
+            out.push(43);
+            out.push(*index);
+        }
         TraceStep::Stloc(index) => {
             out.push(18);
             out.push(*index);
@@ -679,6 +683,7 @@ fn decode_trace_step(cursor: &mut AotCursor<'_>) -> VmResult<TraceStep> {
         15 => TraceStep::Pop,
         16 => TraceStep::Dup,
         17 => TraceStep::Ldloc(cursor.read_u8("ldloc index")?),
+        43 => TraceStep::LdlocCopy(cursor.read_u8("ldloc.copy index")?),
         18 => TraceStep::Stloc(cursor.read_u8("stloc index")?),
         19 => TraceStep::BuiltinCall {
             index: cursor.read_u16("builtin index")?,
@@ -888,6 +893,7 @@ fn validate_aot_trace(trace: &JitTrace, code_len: usize) -> VmResult<()> {
             | TraceStep::Pop
             | TraceStep::Dup
             | TraceStep::Ldloc(_)
+            | TraceStep::LdlocCopy(_)
             | TraceStep::Stloc(_)
             | TraceStep::JumpToRoot
             | TraceStep::Ret => {}
diff --git a/pd-vm/src/vm/jit/native/codegen.rs b/pd-vm/src/vm/jit/native/codegen.rs
@@ -340,6 +340,20 @@ pub(super) fn emit_inline_or_helper_step(
             )?;
             Ok(true)
         }
+        TraceStep::LdlocCopy(index) => {
+            emit_inline_ldloc_copy(
+                b,
+                vm_ptr,
+                helper_ref,
+                exit_block,
+                pointer_type,
+                layout,
+                offsets,
+                *index,
+                root_ip,
+            )?;
+            Ok(true)
+        }
         TraceStep::Stloc(index) => {
             emit_inline_stloc(
                 b,
@@ -2482,6 +2496,7 @@ fn step_to_call(step: &TraceStep, root_ip: usize) -> VmResult<(i64, i64, i64, i6
         TraceStep::Pop => (OP_POP, 0, 0, 0),
         TraceStep::Dup => (OP_DUP, 0, 0, 0),
         TraceStep::Ldloc(index) => (OP_LDLOC, i64::from(*index), 0, 0),
+        TraceStep::LdlocCopy(index) => (OP_LDLOC, i64::from(*index), 0, 0),
         TraceStep::Stloc(index) => (OP_STLOC, i64::from(*index), 0, 0),
         TraceStep::Call {
             index,
diff --git a/pd-vm/src/vm/jit/native/cranelift.rs b/pd-vm/src/vm/jit/native/cranelift.rs
@@ -66,6 +66,9 @@ const OP_JUMP: i64 = 23;
 const OP_BUILTIN_CALL: i64 = 24;
 
 fn fused_ldloc_copy_slot(steps: &[TraceStep], index: usize) -> Option<u8> {
+    if let Some(TraceStep::LdlocCopy(slot)) = steps.get(index) {
+        return Some(*slot);
+    }
     let Some(TraceStep::Ldloc(slot)) = steps.get(index) else {
         return None;
     };
@@ -274,7 +277,12 @@ pub(crate) fn compile_trace(
                     slot,
                     trace.root_ip,
                 )?;
-                step_index += 3;
+                step_index += if matches!(trace.steps.get(step_index), Some(TraceStep::LdlocCopy(_)))
+                {
+                    1
+                } else {
+                    3
+                };
                 continue;
             }
 
diff --git a/pd-vm/src/vm/jit/runtime.rs b/pd-vm/src/vm/jit/runtime.rs
@@ -214,12 +214,7 @@ impl Vm {
             self.native_trace_exec_count
         ));
         if self.jit_runtime_diagnostics_enabled {
-            let mut runtime_entries: Vec<(&'static str, u64)> = self
-                .jit_runtime_counters
-                .iter()
-                .map(|(name, count)| (*name, *count))
-                .collect();
-            runtime_entries.sort_unstable_by_key(|(name, _)| *name);
+            let runtime_entries = self.jit_runtime_diagnostics_snapshot();
             let total_runtime_events = runtime_entries
                 .iter()
                 .fold(0u64, |acc, (_, count)| acc.saturating_add(*count));
@@ -465,6 +460,14 @@ impl Vm {
                     let value = std::mem::replace(slot, crate::bytecode::Value::Null);
                     self.stack.push(value);
                 }
+                TraceStep::LdlocCopy(index) => {
+                    let value = self
+                        .locals
+                        .get(*index as usize)
+                        .cloned()
+                        .ok_or(VmError::InvalidLocal(*index))?;
+                    self.stack.push(value);
+                }
                 TraceStep::Stloc(index) => {
                     let value = self.pop_value()?;
                     self.store_local_with_drop_contract(*index, value)?;
diff --git a/pd-vm/src/vm/jit/trace.rs b/pd-vm/src/vm/jit/trace.rs
@@ -112,6 +112,7 @@ pub enum TraceStep {
     Pop,
     Dup,
     Ldloc(u8),
+    LdlocCopy(u8),
     Stloc(u8),
     BuiltinCall {
         index: u16,
@@ -721,11 +722,7 @@ impl TraceJitEngine {
                 let index = read_u8(code, &mut ip)
                     .ok_or(JitNyiReason::InvalidImmediate("ldloc.copy"))?;
                 step_ips.push(instr_ip);
-                steps.push(TraceStep::Ldloc(index));
-                step_ips.push(instr_ip);
-                steps.push(TraceStep::Dup);
-                step_ips.push(instr_ip);
-                steps.push(TraceStep::Stloc(index));
+                steps.push(TraceStep::LdlocCopy(index));
                 continue;
             }
             if opcode == OpCode::Stloc as u8 {
@@ -987,11 +984,7 @@ impl TraceJitEngine {
                 let index = read_u8(code, &mut ip)
                     .ok_or(JitNyiReason::InvalidImmediate("ldloc.copy"))?;
                 step_ips.push(instr_ip);
-                steps.push(TraceStep::Ldloc(index));
-                step_ips.push(instr_ip);
-                steps.push(TraceStep::Dup);
-                step_ips.push(instr_ip);
-                steps.push(TraceStep::Stloc(index));
+                steps.push(TraceStep::LdlocCopy(index));
                 continue;
             }
             if opcode == OpCode::Stloc as u8 {
@@ -1224,6 +1217,7 @@ fn trace_step_name(step: &TraceStep) -> &'static str {
         TraceStep::Pop => "pop",
         TraceStep::Dup => "dup",
         TraceStep::Ldloc(_) => "ldloc",
+        TraceStep::LdlocCopy(_) => "ldloc_copy",
         TraceStep::Stloc(_) => "stloc",
         TraceStep::BuiltinCall { .. } => "call",
         TraceStep::Call { .. } => "call",
diff --git a/pd-vm/src/vm/mod.rs b/pd-vm/src/vm/mod.rs
diff --git a/results/pd_vm_lua_hot_loop_latest.json b/results/pd_vm_lua_hot_loop_latest.json
diff --git a/results/pd_vm_lua_hot_loop_latest.md b/results/pd_vm_lua_hot_loop_latest.md
