From 734560e6c922afb88f2d4796e611e16a2d7923fa Mon Sep 17 00:00:00 2001 From: Eli Arbel <46826214+eliarbel@users.noreply.github.com> Date: Sun, 8 Sep 2024 15:10:30 +0300 Subject: [PATCH] Oxidize CheckGateDirection (#13042) * First working implementation * A control flow fix, lint, doc and some more polishing * Addressing review comments and further simplifying the code * Address more review comments * Use explicit temp storage for mapping to avoid E0716 in Rust 1.70 * Allow checking 2Q non-control flow instructions --- crates/accelerate/src/gate_direction.rs | 150 ++++++++++++++++++ crates/accelerate/src/lib.rs | 1 + .../accelerate/src/target_transpiler/mod.rs | 2 +- crates/circuit/src/dag_circuit.rs | 30 ++-- crates/circuit/src/lib.rs | 2 +- crates/pyext/src/lib.rs | 7 +- qiskit/__init__.py | 1 + .../passes/utils/check_gate_direction.py | 49 +----- .../transpiler/test_check_gate_direction.py | 26 +-- 9 files changed, 196 insertions(+), 72 deletions(-) create mode 100644 crates/accelerate/src/gate_direction.rs diff --git a/crates/accelerate/src/gate_direction.rs b/crates/accelerate/src/gate_direction.rs new file mode 100644 index 00000000000..aee5ca09732 --- /dev/null +++ b/crates/accelerate/src/gate_direction.rs @@ -0,0 +1,150 @@ +// 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. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +use crate::nlayout::PhysicalQubit; +use crate::target_transpiler::Target; +use hashbrown::HashSet; +use pyo3::prelude::*; +use qiskit_circuit::imports; +use qiskit_circuit::operations::OperationRef; +use qiskit_circuit::{ + dag_circuit::{DAGCircuit, NodeType}, + operations::Operation, + packed_instruction::PackedInstruction, + Qubit, +}; +use smallvec::smallvec; + +/// Check if the two-qubit gates follow the right direction with respect to the coupling map. +/// +/// Args: +/// dag: the DAGCircuit to analyze +/// +/// coupling_edges: set of edge pairs representing a directed coupling map, against which gate directionality is checked +/// +/// Returns: +/// true iff all two-qubit gates comply with the coupling constraints +#[pyfunction] +#[pyo3(name = "check_gate_direction_coupling")] +fn py_check_with_coupling_map( + py: Python, + dag: &DAGCircuit, + coupling_edges: HashSet<[Qubit; 2]>, +) -> PyResult { + let coupling_map_check = + |_: &PackedInstruction, op_args: &[Qubit]| -> bool { coupling_edges.contains(op_args) }; + + check_gate_direction(py, dag, &coupling_map_check, None) +} + +/// Check if the two-qubit gates follow the right direction with respect to instructions supported in the given target. +/// +/// Args: +/// dag: the DAGCircuit to analyze +/// +/// target: the Target against which gate directionality compliance is checked +/// +/// Returns: +/// true iff all two-qubit gates comply with the target's coupling constraints +#[pyfunction] +#[pyo3(name = "check_gate_direction_target")] +fn py_check_with_target(py: Python, dag: &DAGCircuit, target: &Target) -> PyResult { + let target_check = |inst: &PackedInstruction, op_args: &[Qubit]| -> bool { + let qargs = smallvec![ + PhysicalQubit::new(op_args[0].0), + PhysicalQubit::new(op_args[1].0) + ]; + + target.instruction_supported(inst.op.name(), Some(&qargs)) + }; + + check_gate_direction(py, dag, &target_check, None) +} + +// The main routine for checking gate directionality. +// +// gate_complies: a function returning true iff the two-qubit gate direction complies with directionality constraints +// +// qubit_mapping: used for mapping the index of a given qubit within an instruction qargs vector to the corresponding qubit index of the +// original DAGCircuit the pass was called with. This mapping is required since control flow blocks are represented by nested DAGCircuit +// objects whose instruction qubit indices are relative to the parent DAGCircuit they reside in, thus when we recurse into nested DAGs, we need +// to carry the mapping context relative to the original DAG. +// When qubit_mapping is None, the identity mapping is assumed +fn check_gate_direction( + py: Python, + dag: &DAGCircuit, + gate_complies: &T, + qubit_mapping: Option<&[Qubit]>, +) -> PyResult +where + T: Fn(&PackedInstruction, &[Qubit]) -> bool, +{ + for node in dag.op_nodes(false) { + let NodeType::Operation(packed_inst) = &dag.dag[node] else { + panic!("PackedInstruction is expected"); + }; + + let inst_qargs = dag.get_qargs(packed_inst.qubits); + + if let OperationRef::Instruction(py_inst) = packed_inst.op.view() { + if py_inst.control_flow() { + let circuit_to_dag = imports::CIRCUIT_TO_DAG.get_bound(py); // TODO: Take out of the recursion + let py_inst = py_inst.instruction.bind(py); + + for block in py_inst.getattr("blocks")?.iter()? { + let inner_dag: DAGCircuit = circuit_to_dag.call1((block?,))?.extract()?; + + let block_ok = if let Some(mapping) = qubit_mapping { + let mapping = inst_qargs // Create a temp mapping for the recursive call + .iter() + .map(|q| mapping[q.0 as usize]) + .collect::>(); + + check_gate_direction(py, &inner_dag, gate_complies, Some(&mapping))? + } else { + check_gate_direction(py, &inner_dag, gate_complies, Some(inst_qargs))? + }; + + if !block_ok { + return Ok(false); + } + } + continue; + } + } + + if inst_qargs.len() == 2 + && !match qubit_mapping { + // Check gate direction based either on a given custom mapping or the identity mapping + Some(mapping) => gate_complies( + packed_inst, + &[ + mapping[inst_qargs[0].0 as usize], + mapping[inst_qargs[1].0 as usize], + ], + ), + None => gate_complies(packed_inst, inst_qargs), + } + { + return Ok(false); + } + } + + Ok(true) +} + +#[pymodule] +pub fn gate_direction(m: &Bound) -> PyResult<()> { + m.add_wrapped(wrap_pyfunction!(py_check_with_coupling_map))?; + m.add_wrapped(wrap_pyfunction!(py_check_with_target))?; + Ok(()) +} diff --git a/crates/accelerate/src/lib.rs b/crates/accelerate/src/lib.rs index 26d93ab3f65..916aeeb8351 100644 --- a/crates/accelerate/src/lib.rs +++ b/crates/accelerate/src/lib.rs @@ -24,6 +24,7 @@ pub mod edge_collections; pub mod error_map; pub mod euler_one_qubit_decomposer; pub mod filter_op_nodes; +pub mod gate_direction; pub mod inverse_cancellation; pub mod isometry; pub mod nlayout; diff --git a/crates/accelerate/src/target_transpiler/mod.rs b/crates/accelerate/src/target_transpiler/mod.rs index 08b3e90eabb..65fc8e80d75 100644 --- a/crates/accelerate/src/target_transpiler/mod.rs +++ b/crates/accelerate/src/target_transpiler/mod.rs @@ -50,7 +50,7 @@ mod exceptions { } // Custom types -type Qargs = SmallVec<[PhysicalQubit; 2]>; +pub type Qargs = SmallVec<[PhysicalQubit; 2]>; type GateMap = IndexMap; type PropsMap = NullableIndexMap>; type GateMapState = Vec<(String, Vec<(Option, Option)>)>; diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index e57942f87cc..5b9769b4d19 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -3975,19 +3975,12 @@ def _format(operand): } /// Get list of 2 qubit operations. Ignore directives like snapshot and barrier. - fn two_qubit_ops(&self, py: Python) -> PyResult>> { + #[pyo3(name = "two_qubit_ops")] + pub fn py_two_qubit_ops(&self, py: Python) -> PyResult>> { let mut nodes = Vec::new(); - for (node, weight) in self.dag.node_references() { - if let NodeType::Operation(ref packed) = weight { - if packed.op.directive() { - continue; - } - - let qargs = self.qargs_interner.get(packed.qubits); - if qargs.len() == 2 { - nodes.push(self.unpack_into(py, node, weight)?); - } - } + for node in self.two_qubit_ops() { + let weight = self.dag.node_weight(node).expect("NodeIndex in graph"); + nodes.push(self.unpack_into(py, node, weight)?); } Ok(nodes) } @@ -5756,6 +5749,19 @@ impl DAGCircuit { } } + /// Return an iterator of 2 qubit operations. Ignore directives like snapshot and barrier. + pub fn two_qubit_ops(&self) -> impl Iterator + '_ { + Box::new(self.op_nodes(false).filter(|index| { + let weight = self.dag.node_weight(*index).expect("NodeIndex in graph"); + if let NodeType::Operation(ref packed) = weight { + let qargs = self.qargs_interner.get(packed.qubits); + qargs.len() == 2 + } else { + false + } + })) + } + // Filter any nodes that don't match a given predicate function pub fn filter_op_nodes(&mut self, mut predicate: F) where diff --git a/crates/circuit/src/lib.rs b/crates/circuit/src/lib.rs index 4ca86c2ca83..5106ba03028 100644 --- a/crates/circuit/src/lib.rs +++ b/crates/circuit/src/lib.rs @@ -32,7 +32,7 @@ use pyo3::prelude::*; use pyo3::types::{PySequence, PyTuple}; pub type BitType = u32; -#[derive(Copy, Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq)] +#[derive(Copy, Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq, FromPyObject)] pub struct Qubit(pub BitType); #[derive(Copy, Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq)] pub struct Clbit(pub BitType); diff --git a/crates/pyext/src/lib.rs b/crates/pyext/src/lib.rs index 33170380939..96b5e84a9cb 100644 --- a/crates/pyext/src/lib.rs +++ b/crates/pyext/src/lib.rs @@ -17,9 +17,9 @@ use qiskit_accelerate::{ commutation_analysis::commutation_analysis, commutation_checker::commutation_checker, convert_2q_block_matrix::convert_2q_block_matrix, dense_layout::dense_layout, error_map::error_map, euler_one_qubit_decomposer::euler_one_qubit_decomposer, - filter_op_nodes::filter_op_nodes_mod, inverse_cancellation::inverse_cancellation_mod, - isometry::isometry, nlayout::nlayout, optimize_1q_gates::optimize_1q_gates, - pauli_exp_val::pauli_expval, + filter_op_nodes::filter_op_nodes_mod, gate_direction::gate_direction, + inverse_cancellation::inverse_cancellation_mod, isometry::isometry, nlayout::nlayout, + optimize_1q_gates::optimize_1q_gates, pauli_exp_val::pauli_expval, remove_diagonal_gates_before_measure::remove_diagonal_gates_before_measure, results::results, sabre::sabre, sampled_exp_val::sampled_exp_val, sparse_pauli_op::sparse_pauli_op, star_prerouting::star_prerouting, stochastic_swap::stochastic_swap, synthesis::synthesis, @@ -72,6 +72,7 @@ fn _accelerate(m: &Bound) -> PyResult<()> { add_submodule(m, uc_gate, "uc_gate")?; add_submodule(m, utils, "utils")?; add_submodule(m, vf2_layout, "vf2_layout")?; + add_submodule(m, gate_direction, "gate_direction")?; add_submodule(m, commutation_checker, "commutation_checker")?; add_submodule(m, commutation_analysis, "commutation_analysis")?; Ok(()) diff --git a/qiskit/__init__.py b/qiskit/__init__.py index 27830243d85..f731465beec 100644 --- a/qiskit/__init__.py +++ b/qiskit/__init__.py @@ -92,6 +92,7 @@ sys.modules["qiskit._accelerate.commutation_checker"] = _accelerate.commutation_checker sys.modules["qiskit._accelerate.commutation_analysis"] = _accelerate.commutation_analysis sys.modules["qiskit._accelerate.synthesis.linear_phase"] = _accelerate.synthesis.linear_phase +sys.modules["qiskit._accelerate.gate_direction"] = _accelerate.gate_direction sys.modules["qiskit._accelerate.inverse_cancellation"] = _accelerate.inverse_cancellation sys.modules["qiskit._accelerate.check_map"] = _accelerate.check_map sys.modules["qiskit._accelerate.filter_op_nodes"] = _accelerate.filter_op_nodes diff --git a/qiskit/transpiler/passes/utils/check_gate_direction.py b/qiskit/transpiler/passes/utils/check_gate_direction.py index e797be95c4a..5251bfc8de9 100644 --- a/qiskit/transpiler/passes/utils/check_gate_direction.py +++ b/qiskit/transpiler/passes/utils/check_gate_direction.py @@ -12,9 +12,11 @@ """Check if the gates follow the right direction with respect to the coupling map.""" -from qiskit.circuit.controlflow import CONTROL_FLOW_OP_NAMES -from qiskit.converters import circuit_to_dag from qiskit.transpiler.basepasses import AnalysisPass +from qiskit._accelerate.gate_direction import ( + check_gate_direction_coupling, + check_gate_direction_target, +) class CheckGateDirection(AnalysisPass): @@ -34,42 +36,6 @@ def __init__(self, coupling_map, target=None): self.coupling_map = coupling_map self.target = target - def _coupling_map_visit(self, dag, wire_map, edges=None): - if edges is None: - edges = self.coupling_map.get_edges() - # Don't include directives to avoid things like barrier, which are assumed always supported. - for node in dag.op_nodes(include_directives=False): - if node.name in CONTROL_FLOW_OP_NAMES: - for block in node.op.blocks: - inner_wire_map = { - inner: wire_map[outer] for outer, inner in zip(node.qargs, block.qubits) - } - - if not self._coupling_map_visit(circuit_to_dag(block), inner_wire_map, edges): - return False - elif ( - len(node.qargs) == 2 - and (wire_map[node.qargs[0]], wire_map[node.qargs[1]]) not in edges - ): - return False - return True - - def _target_visit(self, dag, wire_map): - # Don't include directives to avoid things like barrier, which are assumed always supported. - for node in dag.op_nodes(include_directives=False): - if node.name in CONTROL_FLOW_OP_NAMES: - for block in node.op.blocks: - inner_wire_map = { - inner: wire_map[outer] for outer, inner in zip(node.qargs, block.qubits) - } - if not self._target_visit(circuit_to_dag(block), inner_wire_map): - return False - elif len(node.qargs) == 2 and not self.target.instruction_supported( - node.name, (wire_map[node.qargs[0]], wire_map[node.qargs[1]]) - ): - return False - return True - def run(self, dag): """Run the CheckGateDirection pass on `dag`. @@ -79,9 +45,8 @@ def run(self, dag): Args: dag (DAGCircuit): DAG to check. """ - wire_map = {bit: i for i, bit in enumerate(dag.qubits)} self.property_set["is_direction_mapped"] = ( - self._coupling_map_visit(dag, wire_map) - if self.target is None - else self._target_visit(dag, wire_map) + check_gate_direction_target(dag, self.target) + if self.target + else check_gate_direction_coupling(dag, set(self.coupling_map.get_edges())) ) diff --git a/test/python/transpiler/test_check_gate_direction.py b/test/python/transpiler/test_check_gate_direction.py index 0743f4c04cf..efd38af1663 100644 --- a/test/python/transpiler/test_check_gate_direction.py +++ b/test/python/transpiler/test_check_gate_direction.py @@ -72,13 +72,13 @@ def test_true_direction(self): def test_true_direction_in_same_layer(self): """Two CXs distance_qubits 1 to each other, in the same layer - qr0:--(+)-- + qr0:---.-- | - qr1:---.--- + qr1:--(+)--- - qr2:--(+)-- + qr2:---.--- | - qr3:---.--- + qr3:--(+)-- CouplingMap map: [0]->[1]->[2]->[3] """ @@ -96,9 +96,9 @@ def test_true_direction_in_same_layer(self): def test_wrongly_mapped(self): """Needs [0]-[1] in a [0]--[2]--[1] - qr0:--(+)-- + qr0:---.--- | - qr1:---.--- + qr1:--(+)-- CouplingMap map: [0]->[2]->[1] """ @@ -115,11 +115,11 @@ def test_wrongly_mapped(self): def test_true_direction_undirected(self): """Mapped but with wrong direction - qr0:--(+)-[H]--.-- + qr0:---.--[H]-(+)- | | - qr1:---.-------|-- + qr1:--(+)------|-- | - qr2:----------(+)- + qr2:-----------.-- CouplingMap map: [1]<-[0]->[2] """ @@ -138,13 +138,13 @@ def test_true_direction_undirected(self): def test_false_direction_in_same_layer_undirected(self): """Two CXs in the same layer, but one is wrongly directed - qr0:--(+)-- + qr0:---.--- | - qr1:---.--- + qr1:--(+)-- - qr2:---.--- + qr2:--(+)-- | - qr3:--(+)-- + qr3:---.--- CouplingMap map: [0]->[1]->[2]->[3] """