feat(jit): specialize Set/ArrayPush on delayed-move collection rebinds

examples/collection_rebind_bench.rs

@@ -40,6 +40,8 @@ impl ExecMode {
 
 struct Measurement {
     samples: Vec<Duration>,
+    trace_attempts: usize,
+    recorded_traces: usize,
     native_traces: usize,
     call_boundary_traces: usize,
     loop_back_traces: usize,
@@ -128,12 +130,14 @@ fn run() -> Result<(), String> {
                 .collect::<Vec<_>>()
                 .join(",");
             println!(
-                "workload={} mode={} width={} iterations={} warmup_runs=1 reused_vm=true median_ns={} native_traces={} call_boundary_traces={} loop_back_traces={} native_execs={} trace_exits={} native_loop_backs={} helper_fallbacks={} generic_builtin_calls={} samples_ns={}",
+                "workload={} mode={} width={} iterations={} warmup_runs=1 reused_vm=true median_ns={} trace_attempts={} recorded_traces={} native_traces={} call_boundary_traces={} loop_back_traces={} native_execs={} trace_exits={} native_loop_backs={} helper_fallbacks={} generic_builtin_calls={} samples_ns={}",
                 workload.label(),
                 mode.label(),
                 config.width,
                 config.iterations,
                 median.as_nanos(),
+                measurement.trace_attempts,
+                measurement.recorded_traces,
                 measurement.native_traces,
                 measurement.call_boundary_traces,
                 measurement.loop_back_traces,
@@ -202,7 +206,9 @@ fn measure(
 
     let snapshot_before = vm.jit_snapshot();
     let metrics_before = snapshot_before.metrics;
-    let native_traces = snapshot_before.traces.len();
+    let trace_attempts = snapshot_before.attempts.len();
+    let recorded_traces = snapshot_before.traces.len();
+    let native_traces = vm.jit_native_trace_count();
     let call_boundary_traces = snapshot_before
         .traces
         .iter()
@@ -217,30 +223,45 @@ fn measure(
     let mut generic_builtin_calls = 0u64;
     for _ in 0..config.samples {
         vm.reset_for_reuse();
+        let native_execs_before_sample = vm.jit_native_exec_count();
         let started = Instant::now();
         let status = vm
             .run()
             .map_err(|err| format!("{} {} run failed: {err}", workload.label(), mode.label()))?;
         let elapsed = started.elapsed();
         verify_result(&vm, status, config)?;
+        if matches!(mode, ExecMode::Jit) && vm.jit_native_exec_count() <= native_execs_before_sample
+        {
+            return Err(format!(
+                "{} measured run did not execute a warmed native trace:\n{}",
+                workload.label(),
+                vm.dump_jit_info()
+            ));
+        }
         generic_builtin_calls = generic_builtin_calls
             .saturating_add(vm.interpreter_metrics_snapshot().generic_builtin_call_count);
         black_box(vm.stack());
         samples.push(elapsed);
     }
     let snapshot_after = vm.jit_snapshot();
-    if snapshot_after.traces.len() != native_traces {
+    let measured_trace_attempts = snapshot_after.attempts.len();
+    let measured_recorded_traces = snapshot_after.traces.len();
+    let measured_native_traces = vm.jit_native_trace_count();
+    if measured_trace_attempts != trace_attempts
+        || measured_recorded_traces != recorded_traces
+        || measured_native_traces != native_traces
+    {
         return Err(format!(
-            "{} {} compiled additional traces during measured runs: warmup={} measured={}",
+            "{} {} changed JIT trace state during measured runs: warmup=attempts:{trace_attempts}/recorded:{recorded_traces}/native:{native_traces} measured=attempts:{measured_trace_attempts}/recorded:{measured_recorded_traces}/native:{measured_native_traces}",
             workload.label(),
             mode.label(),
-            native_traces,
-            snapshot_after.traces.len()
         ));
     }
     let metrics_after = snapshot_after.metrics;
     Ok(Measurement {
         samples,
+        trace_attempts,
+        recorded_traces,
         native_traces,
         call_boundary_traces,
         loop_back_traces,

src/builtins/runtime/core.rs

@@ -220,15 +220,18 @@ fn builtin_array_push_shared_impl(mut items: SharedArray, value: AnyValue) -> Sh
     items
 }
 
-pub(super) fn builtin_array_push(args: &mut [Value]) -> VmResult<CallReturn> {
-    let items = match take_arg(args, 0, "array_push array")? {
+pub(crate) fn builtin_array_push_owned(items: Value, value: Value) -> VmResult<Value> {
+    let items = match items {
         Value::Array(values) => values,
         _ => return Err(VmError::TypeMismatch("array")),
     };
+    Ok(Value::Array(builtin_array_push_shared_impl(items, value)))
+}
+
+pub(super) fn builtin_array_push(args: &mut [Value]) -> VmResult<CallReturn> {
+    let items = take_arg(args, 0, "array_push array")?;
     let value = take_arg(args, 1, "array_push value")?;
-    Ok(return_one(Value::Array(builtin_array_push_shared_impl(
-        items, value,
-    ))))
+    builtin_array_push_owned(items, value).map(return_one)
 }
 
 /// Create an empty map.
@@ -684,20 +687,23 @@ fn builtin_set_map_shared_impl(
     entries
 }
 
-pub(super) fn builtin_set(args: &mut [Value]) -> VmResult<CallReturn> {
-    let container: Value = take_arg(args, 0, "set container")?;
-    let key: Value = take_arg(args, 1, "set key")?;
-    let value: Value = take_arg(args, 2, "set value")?;
+pub(crate) fn builtin_set_owned(container: Value, key: Value, value: Value) -> VmResult<Value> {
     match container {
-        Value::Array(values) => builtin_set_array_shared_impl(values, key.as_int()?, value)
-            .map(|values| return_one(Value::Array(values))),
-        Value::Map(entries) => Ok(return_one(Value::Map(builtin_set_map_shared_impl(
-            entries, key, value,
-        )))),
+        Value::Array(values) => {
+            builtin_set_array_shared_impl(values, key.as_int()?, value).map(Value::Array)
+        }
+        Value::Map(entries) => Ok(Value::Map(builtin_set_map_shared_impl(entries, key, value))),
         _ => Err(VmError::TypeMismatch("array/map")),
     }
 }
 
+pub(super) fn builtin_set(args: &mut [Value]) -> VmResult<CallReturn> {
+    let container = take_arg(args, 0, "set container")?;
+    let key = take_arg(args, 1, "set key")?;
+    let value = take_arg(args, 2, "set value")?;
+    builtin_set_owned(container, key, value).map(return_one)
+}
+
 /// Return an array of container keys or indices.
 #[pd_host_function(name = "keys")]
 pub(super) fn builtin_keys_array_impl(items: VmArrayRef<'_>) -> VmArray {

src/vm/jit/ir.rs

@@ -176,6 +176,15 @@ pub(crate) enum SsaInstKind {
         array: SsaValueId,
         index: SsaValueId,
     },
+    ArraySet {
+        array: SsaValueId,
+        index: SsaValueId,
+        value: SsaValueId,
+    },
+    ArrayPush {
+        array: SsaValueId,
+        value: SsaValueId,
+    },
     MapLen {
         map: SsaValueId,
     },
@@ -187,6 +196,11 @@ pub(crate) enum SsaInstKind {
         map: SsaValueId,
         key: SsaValueId,
     },
+    MapSet {
+        map: SsaValueId,
+        key: SsaValueId,
+        value: SsaValueId,
+    },
     IntNeg {
         input: SsaValueId,
     },
@@ -363,8 +377,15 @@ impl SsaInstKind {
             Self::BytesToArrayU8 { bytes } => vec![*bytes],
             Self::ArrayGet { array, index } => vec![*array, *index],
             Self::ArrayHas { array, index } => vec![*array, *index],
+            Self::ArraySet {
+                array,
+                index,
+                value,
+            } => vec![*array, *index, *value],
+            Self::ArrayPush { array, value } => vec![*array, *value],
             Self::MapGet { map, key } => vec![*map, *key],
             Self::MapHas { map, key } => vec![*map, *key],
+            Self::MapSet { map, key, value } => vec![*map, *key, *value],
             Self::IntAdd { lhs, rhs }
             | Self::IntSub { lhs, rhs }
             | Self::IntMul { lhs, rhs }
@@ -963,9 +984,18 @@ fn render_inst_kind(kind: &SsaInstKind) -> String {
         SsaInstKind::ArrayLen { array } => format!("array_len {array}"),
         SsaInstKind::ArrayGet { array, index } => format!("array_get {array}, {index}"),
         SsaInstKind::ArrayHas { array, index } => format!("array_has {array}, {index}"),
+        SsaInstKind::ArraySet {
+            array,
+            index,
+            value,
+        } => format!("array_set {array}, {index}, {value}"),
+        SsaInstKind::ArrayPush { array, value } => format!("array_push {array}, {value}"),
         SsaInstKind::MapLen { map } => format!("map_len {map}"),
         SsaInstKind::MapGet { map, key } => format!("map_get {map}, {key}"),
         SsaInstKind::MapHas { map, key } => format!("map_has {map}, {key}"),
+        SsaInstKind::MapSet { map, key, value } => {
+            format!("map_set {map}, {key}, {value}")
+        }
         SsaInstKind::IntNeg { input } => format!("ineg {input}"),
         SsaInstKind::IntAdd { lhs, rhs } => format!("iadd {lhs}, {rhs}"),
         SsaInstKind::IntAddImm { lhs, imm } => format!("iadd_imm {lhs}, {imm}"),