diff --git a/quon_na/src/objective.rs b/quon_na/src/objective.rs index e5db35c..73354cd 100644 --- a/quon_na/src/objective.rs +++ b/quon_na/src/objective.rs @@ -641,4 +641,238 @@ mod tests { (spec, trap_transfers, movement_time_us, objective) } } + + #[cfg(feature = "mlir")] + #[test] + fn rydberg_stage_weight_selects_a_different_verified_schedule() { + use crate::entangling_schedule::schedule_entangling_layers; + use crate::geometry::{SpeedModel, SpeedModelKind, movement_duration_for_model}; + use crate::graph::{ + DEFAULT_GAMMA, Interaction, InteractionGraph, InteractionId, InteractionSegment, + LogicalQubitId, SegmentKind, + }; + use crate::layout::{AtomBinding, AtomId, NeutralAtomLayout, SiteId, TrapBinding}; + use crate::lower::{ScheduleLowerParams, lower_schedule}; + use crate::replay::{DeclaredArchitecture, verify_emitted_schedule}; + use crate::schedule::NeutralAtomAction; + use crate::schedule_entry::schedule_from_graph; + use crate::zoned::{ + AwareSearchParams, PlacementCostModel, PlacerMode, ZoneKind, ZoneSpec, + ZonedArchitecture, schedule_zoned_with_aware_params, + }; + + // Two CZ gates already sit on adjacent entanglement pairs. The pairs + // are 10 µm apart, inside the 18.75 µm isolation spacing, so they + // cannot fire together. The other column is 400 µm away and does not + // conflict. One stage pays for that long move. Two stages reuse both + // home pairs and pay for a second Rydberg stage. + let column_pitch_um = 400.0; + let row_pitch_um = 10.0; + let arch = ZonedArchitecture { + zones: vec![ + ZoneSpec { + zone_id: 0, + kind: ZoneKind::Storage, + rows: 2, + cols: 2, + origin_um: (-100.0, -100.0), + site_pitch_um: (4.0, 4.0), + pair_gap_um: None, + }, + ZoneSpec { + zone_id: 1, + kind: ZoneKind::Entanglement, + rows: 2, + cols: 2, + origin_um: (0.0, 0.0), + site_pitch_um: (column_pitch_um, row_pitch_um), + pair_gap_um: Some(6.0), + }, + ], + speed_model: SpeedModel { + kind: SpeedModelKind::Sqrt, + acceleration_m_s2: 2750.0, + jerk_m_s3: 0.0, + max_velocity_m_s: 0.0, + }, + trap_transfer_us: 15, + require_readout_zone: false, + rydberg_range_um: 7.5, + min_rydberg_spacing_um: 18.75, + aod_min_separation_um: 2.0, + }; + // Movement, transfer, and idle stay fixed and positive. Only the + // stage weight changes. The conflicting gate's first legal pair is + // the far column, so both atoms travel hypot(column, row). That is + // one movement group, four trap transfers, and two stationary atoms + // idle for the load, the move, and the store. The verified total + // also charges 1 µs of idle per atom on the extra entangle layer, so + // the low weight sits below that smaller crossover and the high + // weight sits above the placer's cost. + let movement_weight = 1.0; + let transfer_weight = 1.0; + let idle_weight = 1.0; + let travel_um = (column_pitch_um * column_pitch_um + row_pitch_um * row_pitch_um).sqrt(); + let movement_us = movement_duration_for_model(travel_um, &arch.speed_model) as f64; + let trap_us = arch.trap_transfer_us as f64; + let transfers = 4.0; + let stationary_idle_us = 2.0 * (movement_us + 2.0 * trap_us); + let placer_long_move = movement_weight * movement_us + + transfer_weight * transfers + + idle_weight * stationary_idle_us; + let verified_long_move = placer_long_move - idle_weight * 4.0; + let stage_below = verified_long_move - 1.0; + let stage_above = placer_long_move + 1.0; + assert!(stage_below > 0.0); + assert!(stage_below < verified_long_move && verified_long_move < placer_long_move); + assert!(placer_long_move < stage_above); + let below_long_move = weights(stage_below, movement_weight, transfer_weight, idle_weight); + let above_long_move = weights(stage_above, movement_weight, transfer_weight, idle_weight); + + let (split_spec, _, split_obj) = scheduled(&arch, below_long_move); + let (kept_spec, _, kept_obj) = scheduled(&arch, above_long_move); + assert_eq!(split_obj.rydberg_stages, 2); + assert_eq!(kept_obj.rydberg_stages, 1); + assert_eq!(kept_obj.movement_time_us as f64, movement_us); + assert_eq!(kept_obj.trap_transfers, transfers as u64); + assert_eq!(split_obj.movement_time_us, 0); + assert_eq!(split_obj.trap_transfers, 0); + assert_eq!( + split_obj.movement_time_weight, + kept_obj.movement_time_weight + ); + assert_eq!( + split_obj.trap_transfer_weight, + kept_obj.trap_transfer_weight + ); + assert_eq!(split_obj.idle_time_weight, kept_obj.idle_time_weight); + assert!(split_obj.movement_time_weight > 0.0); + assert!(split_obj.trap_transfer_weight > 0.0); + assert!(split_obj.idle_time_weight > 0.0); + let measured_long_move = movement_weight + * (kept_obj.movement_time_us as f64 - split_obj.movement_time_us as f64) + + transfer_weight * (kept_obj.trap_transfers as f64 - split_obj.trap_transfers as f64) + + idle_weight * (kept_obj.idle_time_us as f64 - split_obj.idle_time_us as f64); + assert_eq!(measured_long_move, verified_long_move); + assert!(stage_below < measured_long_move && measured_long_move < stage_above); + + let kept_under_low = + objective_from_verified_schedule(&kept_spec, &below_long_move).expect("rescore"); + let split_under_high = + objective_from_verified_schedule(&split_spec, &above_long_move).expect("rescore"); + assert!(split_obj.total < kept_under_low.total); + assert!(kept_obj.total < split_under_high.total); + + fn scheduled( + arch: &ZonedArchitecture, + weights: NeutralAtomCostModel, + ) -> (crate::dialect::ScheduleSpec, u64, ScheduleObjective) { + let left = InteractionId(0); + let right = InteractionId(1); + let graph = InteractionGraph::from_interactions( + vec![ + LogicalQubitId(0), + LogicalQubitId(1), + LogicalQubitId(2), + LogicalQubitId(3), + ], + vec![ + Interaction { + id: left, + qubits: vec![LogicalQubitId(0), LogicalQubitId(1)], + gate_name: "CZ".into(), + dag_layer: 0, + on_critical_path: false, + }, + Interaction { + id: right, + qubits: vec![LogicalQubitId(2), LogicalQubitId(3)], + gate_name: "CZ".into(), + dag_layer: 0, + on_critical_path: false, + }, + ], + vec![InteractionSegment { + kind: SegmentKind::CommutationGroup, + interactions: vec![left, right], + }], + DEFAULT_GAMMA, + ) + .expect("graph"); + let mut request = + schedule_entangling_layers(schedule_from_graph(graph).expect("request"), 340) + .expect("layers") + .request; + // Storage is sites 0..4. Entanglement pairs are generated row-major: + // (4,5) at y=0, (6,7) at x=400, (8,9) at y=10 beside the first pair. + request.layout = Some(NeutralAtomLayout { + sites: Vec::new(), + initial_bindings: vec![ + AtomBinding { + atom: AtomId(0), + trap: TrapBinding::Slm { site: SiteId(4) }, + }, + AtomBinding { + atom: AtomId(1), + trap: TrapBinding::Slm { site: SiteId(5) }, + }, + AtomBinding { + atom: AtomId(2), + trap: TrapBinding::Slm { site: SiteId(8) }, + }, + AtomBinding { + atom: AtomId(3), + trap: TrapBinding::Slm { site: SiteId(9) }, + }, + ], + declared_initial_bindings: Vec::new(), + }); + let result = schedule_zoned_with_aware_params( + request, + arch, + PlacerMode::RoutingAgnostic, + AwareSearchParams::default(), + PlacementCostModel::Weighted { + weights, + speed_model: arch.speed_model, + trap_transfer_us: arch.trap_transfer_us, + }, + ) + .expect("zoned"); + let layout = result.request.layout.as_ref().expect("layout"); + let spec = lower_schedule( + &result.request, + &ScheduleLowerParams { + target_id: "stage-weight".into(), + rydberg_range_um: arch.rydberg_range_um, + min_rydberg_spacing_um: arch.min_rydberg_spacing_um, + aod_min_separation_um: arch.aod_min_separation_um, + }, + ) + .expect("lower"); + let declared = DeclaredArchitecture { + sites: layout.sites.clone(), + initial_bindings: layout.declared_initial_bindings.clone(), + zones: arch.zones.clone(), + check_zones: true, + require_readout_zone: false, + rydberg_range_um: arch.rydberg_range_um, + min_rydberg_spacing_um: arch.min_rydberg_spacing_um, + aod_min_separation_um: arch.aod_min_separation_um, + }; + verify_emitted_schedule(&spec, &declared).expect("verified schedule"); + let objective = objective_from_verified_schedule(&spec, &weights).expect("objective"); + let movement_time_us = result + .request + .layers + .iter() + .flat_map(|layer| &layer.actions) + .filter_map(|action| match action { + NeutralAtomAction::Move(group) => Some(group.duration_us), + _ => None, + }) + .sum(); + (spec, movement_time_us, objective) + } + } } diff --git a/quon_na/src/zoned.rs b/quon_na/src/zoned.rs index 47ba3f3..334bfb0 100644 --- a/quon_na/src/zoned.rs +++ b/quon_na/src/zoned.rs @@ -261,9 +261,10 @@ pub enum PlacementCostModel { /// architecture_model.md §9, using the target's `cost_model` weights. /// Movement time is `t(d_max)` per group and each moving atom costs two /// trap transfers. An atom that misses a group's load, move, and store - /// layers is idle for `t(d_max) + 2 · trap_transfer_us`. Rydberg stages - /// are charged once per non-empty layer when two full assignments are - /// compared. + /// layers is idle for `t(d_max) + 2 · trap_transfer_us`. A non-empty + /// layer costs one Rydberg stage. When that weight makes a later stage + /// strictly cheaper, the schedule defers the expensive gate instead of + /// keeping every placed gate in this layer. Weighted { weights: backend::NeutralAtomCostModel, speed_model: SpeedModel, @@ -1074,6 +1075,11 @@ pub fn schedule_zoned_with_aware_params( assign_aware_legal(&gate_atoms, &inputs, &aware_search, &layout) } }; + // Stage weight cannot change a choice inside one layer: every + // non-empty assignment pays for exactly one Rydberg stage. Spend a + // later stage when that is strictly cheaper than the simultaneous move. + let assignment = + split_weighted_layer_on_stage_cost(assignment, &gate_atoms, &inputs, &layout); aware_search_node_expansions += assignment.node_expansions as u64; match assignment.outcome { AwareSearchOutcome::NotApplicable => {} @@ -1864,6 +1870,7 @@ pub enum AwareSearchOutcome { /// Result of a layer's pair assignment: gates placed this stage (with their /// oriented pair positions) and gates deferred to a follow-up stage. +#[derive(Clone)] struct GateAssignment { placed: Vec<(usize, (Position, Position))>, deferred: Vec, @@ -2382,6 +2389,164 @@ fn assignment_weighted_cost( movement + weights.rydberg_stage_weight * stages as f64 } +/// Defer a placed gate when a later Rydberg stage is strictly cheaper. +/// +/// `assignment_weighted_cost` charges one stage whenever `placed` is +/// non-empty, so the stage weight never separates two assignments that both +/// entangle this layer. This pass keeps the assignment unless moving one +/// gate into the follow-up layer lowers the §9 cost. `Time` and +/// `ErrorBudget` are unchanged. Equal costs keep the original schedule. +fn split_weighted_layer_on_stage_cost( + mut assignment: GateAssignment, + gates: &[(AtomId, AtomId)], + inputs: &AssignInputs<'_>, + layout: &NeutralAtomLayout, +) -> GateAssignment { + if !matches!(inputs.cost_model, PlacementCostModel::Weighted { .. }) { + return assignment; + } + loop { + if assignment.placed.len() < 2 { + break; + } + let Some(full_follow) = followup_layer_cost( + &assignment.deferred, + &assignment.placed, + gates, + inputs, + layout, + ) else { + break; + }; + let full = layer_weighted_cost(&assignment, gates, inputs, layout) + full_follow; + let mut best: Option<(f64, GateAssignment)> = None; + for drop_at in 0..assignment.placed.len() { + let mut kept = assignment.clone(); + let (gate_index, _) = kept.placed.remove(drop_at); + let mut later = kept.deferred.clone(); + later.push(gate_index); + let Some(later_cost) = followup_layer_cost(&later, &kept.placed, gates, inputs, layout) + else { + continue; + }; + let split = layer_weighted_cost(&kept, gates, inputs, layout) + later_cost; + if split < full && best.as_ref().is_none_or(|(cost, _)| split < *cost) { + kept.deferred.push(gate_index); + best = Some((split, kept)); + } + } + match best { + Some((_, next)) => assignment = next, + None => break, + } + } + assignment +} + +fn layer_weighted_cost( + assignment: &GateAssignment, + gates: &[(AtomId, AtomId)], + inputs: &AssignInputs<'_>, + layout: &NeutralAtomLayout, +) -> f64 { + let PlacementCostModel::Weighted { + weights, + speed_model, + trap_transfer_us, + } = inputs.cost_model + else { + return 0.0; + }; + assignment_weighted_cost( + assignment, + gates, + inputs.atom_pos, + layout, + &weights, + &speed_model, + trap_transfer_us, + inputs.aod_min_sep_um, + ) +} + +/// Cost of scheduling `gate_indices` as one later layer, after `placed` has +/// moved. `None` when that layer cannot place every gate (the follow-up +/// would itself defer, so this pass will not pretend it is one stage). +fn followup_layer_cost( + gate_indices: &[usize], + placed: &[(usize, (Position, Position))], + gates: &[(AtomId, AtomId)], + inputs: &AssignInputs<'_>, + layout: &NeutralAtomLayout, +) -> Option { + if gate_indices.is_empty() { + return Some(0.0); + } + let PlacementCostModel::Weighted { + weights, + speed_model, + trap_transfer_us, + } = inputs.cost_model + else { + return None; + }; + let later_gates: Vec<(AtomId, AtomId)> = + gate_indices.iter().map(|&index| gates[index]).collect(); + let (atom_pos, site_occupant) = + positions_after_placed(placed, gates, inputs.atom_pos, inputs.site_occupant, layout); + let later_inputs = AssignInputs { + atom_pos: &atom_pos, + pairs: inputs.pairs, + pair_sites: inputs.pair_sites, + site_occupant: &site_occupant, + conflict_um: inputs.conflict_um, + aod_min_sep_um: inputs.aod_min_sep_um, + cost_model: inputs.cost_model, + }; + let later = assign_greedy_legal(&later_gates, &later_inputs); + if later.placed.len() != later_gates.len() || !later.deferred.is_empty() { + return None; + } + Some(assignment_weighted_cost( + &later, + &later_gates, + &atom_pos, + layout, + &weights, + &speed_model, + trap_transfer_us, + inputs.aod_min_sep_um, + )) +} + +fn positions_after_placed( + placed: &[(usize, (Position, Position))], + gates: &[(AtomId, AtomId)], + atom_pos: &BTreeMap, + site_occupant: &BTreeMap, + layout: &NeutralAtomLayout, +) -> (BTreeMap, BTreeMap) { + let mut atom_pos = atom_pos.clone(); + let mut site_occupant = site_occupant.clone(); + for &(gate_index, (pa, pb)) in placed { + let (a, b) = gates[gate_index]; + for (atom, target) in [(a, pa), (b, pb)] { + let cur = atom_pos.get(&atom).copied().unwrap_or(target); + if euclidean_um(cur, target) < 1e-9 { + continue; + } + let from_site = nearest_site_id(layout, cur); + let to_site = nearest_site_id(layout, target); + atom_pos.insert(atom, target); + if site_occupant.get(&from_site) == Some(&atom) { + site_occupant.remove(&from_site); + } + site_occupant.insert(to_site, atom); + } + } + (atom_pos, site_occupant) +} + /// Choose the routing-agnostic layer assignment: compute both the /// min-weight matching placer (assign_matching_legal) and the greedy /// placer (assign_greedy_legal), keep whichever yields fewer AOD movement