From 67b2555ef65ab4c2b920a0b08893df161a8e2125 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 1 Oct 2026 17:35:01 +0000 Subject: [PATCH] [na] Echo scored cost_model weights in stats The zoned placer scores with the target cost_model. Stats now echo those four weights next to the Time or ErrorBudget mode. Flat AOD does not score with the weights, so that echo stays absent. Refs #484 Co-authored-by: Arnab Ghosh --- quon_na/src/lib.rs | 4 +- quon_na/src/pipeline.rs | 1 + quon_na/src/plan.rs | 24 +++++++- quon_na/src/qec_schedule.rs | 1 + quon_na/src/stats.rs | 37 +++++++++++- ...ats_snapshots__zoned_routing_agnostic.snap | 6 ++ ...hots__zoned_routing_aware_config_only.snap | 6 ++ quon_na/tests/stats_snapshots.rs | 60 ++++++++++++++++++- 8 files changed, 132 insertions(+), 7 deletions(-) diff --git a/quon_na/src/lib.rs b/quon_na/src/lib.rs index e3c13b78..6ae82440 100644 --- a/quon_na/src/lib.rs +++ b/quon_na/src/lib.rs @@ -130,8 +130,8 @@ pub use schedule_view::{ ScheduleViewZone, }; pub use stats::{ - CompactionConfig, EffectiveConfig, NA_STATS_KIND, NA_STATS_SCHEMA_VERSION, NaStats, - SearchDiagnostics, StageTimingsUs, VersionInfo, na_stats_to_json, + CompactionConfig, CostModelWeights, EffectiveConfig, NA_STATS_KIND, NA_STATS_SCHEMA_VERSION, + NaStats, SearchDiagnostics, StageTimingsUs, VersionInfo, na_stats_to_json, }; pub use zoned::{ AWARE_NODE_BUDGET, AgnosticPlacerMechanism, AwareSearchOutcome, AwareSearchParams, diff --git a/quon_na/src/pipeline.rs b/quon_na/src/pipeline.rs index 4a718047..3426891f 100644 --- a/quon_na/src/pipeline.rs +++ b/quon_na/src/pipeline.rs @@ -807,6 +807,7 @@ fn finish_pipeline( placement_strategy, compaction: compaction_config, objective: opts.objective, + cost_model: backend_info.cost_model, }, stage_timings_us: StageTimingsUs { extract_us: None, diff --git a/quon_na/src/plan.rs b/quon_na/src/plan.rs index 7ca21709..976c6f98 100644 --- a/quon_na/src/plan.rs +++ b/quon_na/src/plan.rs @@ -26,7 +26,7 @@ use crate::pipeline::{ }; use crate::placement::{PlacementStrategy, place}; use crate::schedule_entry::GraphScheduleRequest; -use crate::stats::SearchDiagnostics; +use crate::stats::{CostModelWeights, SearchDiagnostics}; use crate::zoned::{ AgnosticPlacerMechanism, PlacementCostModel, PlacerMode, schedule_zoned_with_aware_params, }; @@ -65,6 +65,9 @@ pub struct BackendStageInfo { /// `Some(ScheduleOptimality::Exact)` when z3 proved optimality; /// `Some(ScheduleOptimality::Heuristic)` on timeout/fallback. pub schedule_optimality: Option, + /// §9 weights passed to the zoned placer. `None` for flat AOD, which + /// does not score with `cost_model`. + pub cost_model: Option, } /// Run the place → AOD movement (or zoned routing) backend on a @@ -130,6 +133,12 @@ pub fn plan_backend( ) } }; + info.cost_model = match cost_model { + PlacementCostModel::Weighted { weights, .. } => { + Some(CostModelWeights::from(weights)) + } + PlacementCostModel::Time | PlacementCostModel::ErrorBudget { .. } => None, + }; let stage_started = Instant::now(); let zoned = schedule_zoned_with_aware_params( req, @@ -271,6 +280,8 @@ pub struct QecStageAccumulator { pub aware_search_status: Option<(u64, u64)>, pub agnostic_placer_mechanism: Option, pub schedule_optimality: Option, + /// §9 weights the zoned placer scored with. Same across phases. + pub cost_model: Option, /// Per-phase exact state-prep solver outcome, aggregated across CNOT /// phases (issue #397). `None` when the heuristic scheduler ran (default /// `state_prep` mode, or a per-phase empty CNOT list). `Timeout` is @@ -305,6 +316,9 @@ impl QecStageAccumulator { if self.placement_strategy.is_none() { self.placement_strategy = backend.placement_strategy; } + if self.cost_model.is_none() { + self.cost_model = backend.cost_model; + } // search_diagnostics: accumulate sums. let acc = &mut self.search_diagnostics; let inc = &backend.search_diagnostics; @@ -441,9 +455,16 @@ mod tests { #[test] fn accumulator_accumulates() { let mut acc = QecStageAccumulator::default(); + let weights = CostModelWeights { + rydberg_stage_weight: 7.0, + movement_time_weight: 0.25, + trap_transfer_weight: 3.0, + idle_time_weight: 0.5, + }; let backend = BackendStageInfo { placement_us: Some(10), movement_us: Some(20), + cost_model: Some(weights), ..Default::default() }; acc.accumulate_phase(5, 8, &backend); @@ -451,6 +472,7 @@ mod tests { assert_eq!(acc.entangling_layers_us, 8); assert_eq!(acc.placement_us, Some(10)); assert_eq!(acc.movement_us, Some(20)); + assert_eq!(acc.cost_model, Some(weights)); let backend2 = BackendStageInfo { movement_us: Some(30), diff --git a/quon_na/src/qec_schedule.rs b/quon_na/src/qec_schedule.rs index e10adb9a..91dace97 100644 --- a/quon_na/src/qec_schedule.rs +++ b/quon_na/src/qec_schedule.rs @@ -255,6 +255,7 @@ fn schedule_expanded( placement_strategy: stage_acc.placement_strategy, compaction: compaction_config, objective: opts.objective, + cost_model: stage_acc.cost_model, }, stage_timings_us: StageTimingsUs { extract_us: None, diff --git a/quon_na/src/stats.rs b/quon_na/src/stats.rs index 530f1134..fe21e23d 100644 --- a/quon_na/src/stats.rs +++ b/quon_na/src/stats.rs @@ -186,6 +186,35 @@ pub struct CompactionConfig { pub legality_checked: bool, } +/// §9 `cost_model` weights the zoned placer scored with. +/// +/// Copied from the [`crate::zoned::PlacementCostModel::Weighted`] vector +/// passed to `schedule_zoned_with_aware_params`. Flat AOD does not score +/// with these weights, so the echo stays absent there. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CostModelWeights { + #[serde(default)] + pub rydberg_stage_weight: f64, + #[serde(default)] + pub movement_time_weight: f64, + #[serde(default)] + pub trap_transfer_weight: f64, + #[serde(default)] + pub idle_time_weight: f64, +} + +impl From for CostModelWeights { + fn from(weights: backend::NeutralAtomCostModel) -> Self { + Self { + rydberg_stage_weight: weights.rydberg_stage_weight, + movement_time_weight: weights.movement_time_weight, + trap_transfer_weight: weights.trap_transfer_weight, + idle_time_weight: weights.idle_time_weight, + } + } +} + /// Effective (as-run, not merely as-requested) pipeline configuration. #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -202,10 +231,14 @@ pub struct EffectiveConfig { pub placement_strategy: Option, #[serde(default)] pub compaction: CompactionConfig, - /// Placement/routing objective (issue #309): `Time` (default) or - /// `ErrorBudget`. Mirrors the `--na-objective` CLI knob. + /// Placement/routing objective mode: `Time` (default) or `ErrorBudget`. + /// Mirrors the `--na-objective` CLI knob. The weights that mode scored + /// with are [`Self::cost_model`]. #[serde(default)] pub objective: NaObjective, + /// §9 weights the zoned optimizer scored with. `None` for flat AOD. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cost_model: Option, } /// Tool/target identifiers so a stats file is self-describing without a diff --git a/quon_na/tests/snapshots/stats_snapshots__zoned_routing_agnostic.snap b/quon_na/tests/snapshots/stats_snapshots__zoned_routing_agnostic.snap index d15b09f2..7423a5cb 100644 --- a/quon_na/tests/snapshots/stats_snapshots__zoned_routing_agnostic.snap +++ b/quon_na/tests/snapshots/stats_snapshots__zoned_routing_agnostic.snap @@ -11,6 +11,12 @@ expression: redact_timings(&json) "legality_checked": false, "requested": true }, + "cost_model": { + "idle_time_weight": 1e-6, + "movement_time_weight": 1.0, + "rydberg_stage_weight": 1.0, + "trap_transfer_weight": 1.0 + }, "objective": "time", "placer_mode": "routing_agnostic" }, diff --git a/quon_na/tests/snapshots/stats_snapshots__zoned_routing_aware_config_only.snap b/quon_na/tests/snapshots/stats_snapshots__zoned_routing_aware_config_only.snap index 485c51bf..94141e3b 100644 --- a/quon_na/tests/snapshots/stats_snapshots__zoned_routing_aware_config_only.snap +++ b/quon_na/tests/snapshots/stats_snapshots__zoned_routing_aware_config_only.snap @@ -11,6 +11,12 @@ expression: redact_timings(&json) "legality_checked": false, "requested": true }, + "cost_model": { + "idle_time_weight": 1e-6, + "movement_time_weight": 1.0, + "rydberg_stage_weight": 1.0, + "trap_transfer_weight": 1.0 + }, "objective": "time", "placer_mode": "routing_aware" }, diff --git a/quon_na/tests/stats_snapshots.rs b/quon_na/tests/stats_snapshots.rs index 3bec636a..1a2de6a9 100644 --- a/quon_na/tests/stats_snapshots.rs +++ b/quon_na/tests/stats_snapshots.rs @@ -18,8 +18,8 @@ use std::path::PathBuf; use quon_na::{ - NaBackendKind, NaScheduleOptions, NaStats, PlacementStrategy, PlacerMode, - cubic_commutation_graph, na_stats_to_json, run_from_graph, + CostModelWeights, NaBackendKind, NaObjective, NaScheduleOptions, NaStats, PlacementStrategy, + PlacerMode, cubic_commutation_graph, na_stats_to_json, run_from_graph, }; use serde_json::Value; @@ -188,6 +188,10 @@ fn flat_aod_stats_shape() { stats.config.placement_strategy, Some(PlacementStrategy::RowMajor) ); + assert!( + stats.config.cost_model.is_none(), + "flat AOD does not score with cost_model weights" + ); // Aware-search concept doesn't apply to flat AOD. assert_eq!(stats.search.aware_search_completed_layers, None); assert_eq!(stats.search.aware_search_node_expansions, None); @@ -223,6 +227,58 @@ fn no_compact_flag_reflects_requested_not_applied() { assert!(stats.stage_timings_us.compaction_us.is_none()); } +/// Zoned stats echo the §9 weights the placer scored with, not only the +/// `Time` / `ErrorBudget` mode name. +#[test] +fn stats_echo_scored_cost_weights_not_only_the_mode() { + let mut na = na_target(); + let weights = backend::NeutralAtomCostModel { + rydberg_stage_weight: 7.0, + movement_time_weight: 0.25, + trap_transfer_weight: 3.0, + idle_time_weight: 0.5, + }; + na.cost_model = weights; + assert!( + na.error_model.is_some(), + "error-budget mode needs the target error_model" + ); + + for objective in [NaObjective::Time, NaObjective::ErrorBudget] { + let artifacts = run_from_graph( + cubic_commutation_graph(4).expect("cubic"), + &na, + NaScheduleOptions { + backend: NaBackendKind::Zoned, + placer: PlacerMode::RoutingAgnostic, + objective, + ..Default::default() + }, + None, + ) + .expect("schedule"); + let stats = artifacts + .stats + .expect("run_from_graph always populates stats"); + assert_eq!(stats.config.objective, objective); + assert_eq!( + stats.config.cost_model, + Some(CostModelWeights::from(weights)), + "echoed weights must be the vector the zoned placer scored with" + ); + + let json = na_stats_to_json(&stats).expect("serialize NaStats"); + let value: Value = serde_json::from_str(&json).expect("parse NaStats JSON"); + let mode = serde_json::to_value(objective).expect("mode name"); + assert_eq!(value["config"]["objective"], mode); + let echoed = &value["config"]["cost_model"]; + assert_eq!(echoed["rydberg_stage_weight"], 7.0); + assert_eq!(echoed["movement_time_weight"], 0.25); + assert_eq!(echoed["trap_transfer_weight"], 3.0); + assert_eq!(echoed["idle_time_weight"], 0.5); + } +} + /// A stats JSON with only the fields present in schema v1 must still /// deserialize (the `#[serde(default)]` field-evolution policy, issue #307). #[test]