diff --git a/crates/transpiler/src/passes/sabre/dag.rs b/crates/transpiler/src/passes/sabre/dag.rs index 0343e98fe59e..c234ee2c6716 100644 --- a/crates/transpiler/src/passes/sabre/dag.rs +++ b/crates/transpiler/src/passes/sabre/dag.rs @@ -37,17 +37,23 @@ pub enum InteractionKind { /// the Sabre version and the DAG representation of the block, so we can reconstruct things /// later. When control-flow ops are represented with native DAGs, we won't need to store the /// temporary. - ControlFlow(Box<[(SabreDAG, DAGCircuit)]>), + ControlFlow(Box<[(SabreDAG, DAGCircuit)]>, Box<[VirtualQubit]>), } impl InteractionKind { - fn from_control_flow(cf: ControlFlowView) -> Result { - let blocks: Box<[_]> = cf - .blocks() - .into_iter() - .map(|dag| Ok((SabreDAG::from_dag(dag)?, dag.clone()))) - .collect::>()?; - Ok(Self::ControlFlow(blocks)) - } + fn from_control_flow( + cf: ControlFlowView, + outer_qargs: &[Qubit], +) -> Result { + let blocks: Box<[_]> = cf + .blocks() + .into_iter() + .map(|dag| Ok((SabreDAG::from_dag(dag)?, dag.clone()))) + .collect::>()?; + let outer_virtual: Box<[VirtualQubit]> = outer_qargs + .iter() + .map(|q| VirtualQubit::new(q.0)) + .collect(); + Ok(Self::ControlFlow(blocks, outer_virtual))} fn from_op(op: &PackedOperation, qargs: &[Qubit]) -> Result { if op.directive() { @@ -159,7 +165,7 @@ impl SabreDAG { panic!("op nodes should always be of type `Operation`"); }; let kind = if let Some(cf) = dag.try_view_control_flow(inst) { - InteractionKind::from_control_flow(cf)? + InteractionKind::from_control_flow(cf, dag.get_qargs(inst.qubits))? } else { InteractionKind::from_op(&inst.op, dag.get_qargs(inst.qubits))? }; @@ -234,7 +240,9 @@ impl SabreDAG { for node in self.dag.node_weights() { let kind = match &node.kind { InteractionKind::Synchronize | InteractionKind::TwoQ(_) => node.kind.clone(), - InteractionKind::ControlFlow(_) => InteractionKind::Synchronize, + InteractionKind::ControlFlow(blocks, outer_qargs) => { + extract_dominant_interaction_with_mapping(blocks, outer_qargs) + }, }; // `NodeWeights` guarantees that the weights come out in index order, so the new indexes // must match those in `self`. @@ -260,3 +268,33 @@ impl SabreDAG { out_dag } } + + +fn extract_dominant_interaction_with_mapping( + blocks: &[(SabreDAG, DAGCircuit)], + outer_qargs: &[VirtualQubit], +) -> InteractionKind { + let mut pair_counts: hashbrown::HashMap<[u32; 2], usize> = + hashbrown::HashMap::new(); + let mut best: Option<([VirtualQubit; 2], usize)> = None; + + for (sabre_block, _) in blocks.iter() { + for node in sabre_block.dag.node_weights() { + if let InteractionKind::TwoQ([a, b]) = &node.kind { + let outer_a = outer_qargs.get(a.0 as usize).copied(); + let outer_b = outer_qargs.get(b.0 as usize).copied(); + if let (Some(oa), Some(ob)) = (outer_a, outer_b) { + let mut key = [oa.0, ob.0]; + key.sort_unstable(); + let count = pair_counts.entry(key).or_insert(0); + *count += 1; + if best.map_or(true, |(_, c)| *count > c) { + best = Some(([oa, ob], *count)); + } + } + } + } + } + best.map(|([a, b], _)| InteractionKind::TwoQ([a, b])) + .unwrap_or(InteractionKind::Synchronize) +} \ No newline at end of file diff --git a/crates/transpiler/src/passes/sabre/route.rs b/crates/transpiler/src/passes/sabre/route.rs index 879b4c0c7c53..35a11bdc6a4d 100644 --- a/crates/transpiler/src/passes/sabre/route.rs +++ b/crates/transpiler/src/passes/sabre/route.rs @@ -526,7 +526,7 @@ impl State { continue; } } - InteractionKind::ControlFlow(blocks) => { + InteractionKind::ControlFlow(blocks, _outer_qargs) => { let dag_node_id = *node.indices.first().expect( "if control-flow interactions are included, so are original DAG indices", ); diff --git a/test/python/transpiler/test_sabre_control_flow.py b/test/python/transpiler/test_sabre_control_flow.py new file mode 100644 index 000000000000..85846be3fe1f --- /dev/null +++ b/test/python/transpiler/test_sabre_control_flow.py @@ -0,0 +1,53 @@ +# This code is part of Qiskit. +# +# (C) Copyright IBM 2024 +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. + +"""Test SABRE layout with control flow operations.""" + +import unittest +from qiskit import QuantumCircuit +from qiskit.transpiler import CouplingMap +from qiskit.providers.fake_provider import GenericBackendV2 +from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager +from test import QiskitTestCase + + +class TestSabreControlFlow(QiskitTestCase): + def setUp(self): + super().setUp() + cmap = CouplingMap.from_heavy_hex(5) + self.backend = GenericBackendV2( + cmap.size(), + control_flow=True, + coupling_map=cmap, + seed=42, + ) + + def test_if_else_layout_improvement(self): + pm = generate_preset_pass_manager( + optimization_level=1, + backend=self.backend, + seed_transpiler=42, + ) + qc = QuantumCircuit(10, 1) + for i in range(0, 9, 2): + qc.cx(i, i + 1) + for i in range(1, 8, 2): + qc.cx(i, i + 1) + qc.measure(0, 0) + with qc.if_test((0, 1)): + for i in range(0, 9, 2): + qc.swap(i, i + 1) + for i in range(1, 8, 2): + qc.swap(i, i + 1) + result = pm.run(qc) + depth = result.depth(lambda g: g.operation.num_qubits == 2) + self.assertLess(depth, 20) + + +if __name__ == "__main__": + unittest.main()