DRAFT: [CQT-13] Implement two-qubit (generic) merge pass

opensquirrel/circuit_builder.py

@@ -212,6 +212,10 @@ def add_instruction(self, instruction: Instruction | Iterable[Instruction]) -> S
             self.ir.add_statement(instr)
         return self
 
+    def clear(self) -> None:
+        """Clears all statements from the IR."""
+        self.ir.clear()
+
     def to_circuit(self) -> Circuit:
         """Build the circuit.
 

opensquirrel/ir/ir.py

@@ -162,3 +162,7 @@ def __eq__(self, other: object) -> bool:
 
     def __repr__(self) -> str:
         return f"IR: {self.statements}"
+
+    def clear(self) -> None:
+        """Clears all statements from the IR."""
+        self.statements.clear()

opensquirrel/passes/merger/two_qubit_gates_merger.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import networkx as nx
+import numpy as np
+
+from opensquirrel.circuit_builder import CircuitBuilder
+from opensquirrel.circuit_matrix_calculator import get_circuit_matrix
+from opensquirrel.ir import IR, Gate, Instruction, Qubit
+from opensquirrel.ir.semantics.matrix_gate import MatrixGateSemantic
+from opensquirrel.ir.single_qubit_gate import SingleQubitGate
+from opensquirrel.ir.two_qubit_gate import TwoQubitGate
+from opensquirrel.passes.merger.general_merger import Merger
+
+if TYPE_CHECKING:
+    from opensquirrel.circuit import Circuit
+
+
+def build_graph(ir: IR) -> nx.DiGraph:
+    n = len(ir.statements)
+    graph = nx.DiGraph()
+    graph.add_nodes_from(
+        (i, {"qubit_indices": statement.qubit_indices})
+        for i, statement in enumerate(ir.statements)
+        if isinstance(statement, Instruction)
+    )
+
+    for i, statement in enumerate(ir.statements):
+        if not isinstance(statement, Instruction):
+            continue
+
+        qubit_indices = set(statement.qubit_indices)
+        for j in range(i + 1, n):
+            other_statement = ir.statements[j]
+            if isinstance(other_statement, Instruction):
+                other_qubit_indices = set(other_statement.qubit_indices)
+
+                if inter := qubit_indices.intersection(other_qubit_indices):
+                    graph.add_edge(i, j, qubit_index=tuple(inter))
+                    qubit_indices = qubit_indices.difference(inter)
+
+            if not qubit_indices:
+                break
+
+    if not nx.is_directed_acyclic_graph(graph):
+        raise ValueError
+
+    return graph
+
+
+def group_gates(graph: nx.DiGraph) -> list[tuple[set[int], set[int]]]:
+    groups: list[tuple[set, set]] = []
+    available_nodes = set(graph.nodes)
+
+    if len(available_nodes) == 1:
+        return [(available_nodes, set(graph.nodes[0]["qubit_indices"]))]
+
+    for edge in graph.edges():
+        source_node = graph.nodes[edge[0]]
+        target_node = graph.nodes[edge[1]]
+
+        edge_nodes = {edge[0], edge[1]}
+        active_nodes = edge_nodes & available_nodes
+        qubit_indices = set(source_node["qubit_indices"]) | set(target_node["qubit_indices"])
+
+        matched_group = False
+        for group, indices in groups:
+            if not (group & edge_nodes):
+                continue
+
+            if len(indices) == 1:
+                indices.update(qubit_indices)
+
+            if qubit_indices.issubset(indices):
+                group.update(active_nodes)
+                available_nodes.difference_update(active_nodes)
+
+            matched_group = True
+            break
+
+        if not matched_group:
+            groups.append((active_nodes, set(qubit_indices)))
+            available_nodes.difference_update(active_nodes)
+
+    # Sort the groups, such that the order of the two qubit gates is preserved.
+    return sorted(groups, key=lambda x: _first_two_qubit_gate(graph, x[0]))
+
+
+def _first_two_qubit_gate(graph: nx.DiGraph, group: set[int]) -> int:
+    """Return the first statement index pointing to a two qubit gate."""
+    return min(i for i in group if len(graph.nodes[i]["qubit_indices"]) == 2)
+
+
+def normalize_gate_indices(gate: Gate) -> Gate:
+    if isinstance(gate, TwoQubitGate):
+        gate.qubit0 = Qubit(gate.qubit0.index % 2)
+        gate.qubit1 = Qubit(gate.qubit1.index % 2)
+        return gate
+
+    if isinstance(gate, SingleQubitGate):
+        gate.qubit = Qubit(gate.qubit.index % 2)
+        return gate
+
+    msg = f"Unsupported gate type: {type(gate)}"
+    raise TypeError(msg)
+
+
+def _merge_gate_group(ir: IR, group: set[int], qubit_indices: set[int]) -> TwoQubitGate:
+    builder = CircuitBuilder(len(qubit_indices))
+    for index in group:
+        statement = ir.statements[index]
+        if isinstance(statement, Gate):
+            statement = normalize_gate_indices(statement)
+            builder.add_instruction(statement)
+
+    sub_circuit_matrix = _get_sub_circuit_matrix(builder.to_circuit())
+    return TwoQubitGate(*qubit_indices, gate_semantic=MatrixGateSemantic(sub_circuit_matrix))
+
+
+def _get_sub_circuit_matrix(circuit: Circuit) -> np.ndarray:
+    # `get_circuit_matrix` uses the convention of the first qubit being the most significant bit,
+    # so we need to swap the qubits before and after calculating the matrix
+    swap = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
+    return swap @ get_circuit_matrix(circuit) @ swap
+
+
+class TwoQubitGatesMerger(Merger):
+    def merge(self, ir: IR, qubit_register_size: int) -> None:
+        """Merge all consecutive two-qubit gates in the circuit.
+
+        Args:
+            ir (IR): Intermediate representation of the circuit.
+            qubit_register_size (int): Size of the qubit register
+
+        """
+        graph = build_graph(ir)
+        if len(graph.nodes) == 1:
+            return
+
+        groups = group_gates(graph)
+        ir.statements = [_merge_gate_group(ir, group, qubit_indices) for group, qubit_indices in groups]

tests/passes/merger/test_two_qubit_gates_merger.py

@@ -0,0 +1,131 @@
+import pytest
+
+from opensquirrel import CNOT, H, circuit_matrix_calculator
+from opensquirrel.circuit import Circuit
+from opensquirrel.circuit_builder import CircuitBuilder
+from opensquirrel.common import are_matrices_equivalent_up_to_global_phase
+from opensquirrel.ir import IR, Statement
+from opensquirrel.passes.merger.two_qubit_gates_merger import TwoQubitGatesMerger, build_graph, group_gates
+
+
+class TestNodeGraph:
+    @pytest.mark.parametrize(
+        ("statements", "expected_edges"),
+        [
+            (
+                [H(0), H(1), H(2)],
+                set(),
+            ),
+            (
+                [H(0), CNOT(0, 1), H(1)],
+                {(0, 1), (1, 2)},
+            ),
+            (
+                [H(1), CNOT(0, 1), H(1), H(2), CNOT(1, 2), H(2)],
+                {(0, 1), (1, 2), (2, 4), (3, 4), (4, 5)},
+            ),
+            (
+                [H(3), CNOT(2, 3), H(1), H(2), CNOT(0, 1), H(2), H(1), H(3), CNOT(1, 2), H(2)],
+                {(0, 1), (1, 3), (1, 7), (6, 8), (2, 4), (4, 6), (3, 5), (5, 8), (8, 9)},
+            ),
+        ],
+    )
+    def test_build_graph_creates_dependencies(
+        self, statements: list[Statement], expected_edges: set[tuple[int, int]]
+    ) -> None:
+        ir = IR()
+        for statement in statements:
+            ir.add_statement(statement)
+
+        graph = build_graph(ir)
+
+        assert set(graph.edges()) == expected_edges
+
+    @pytest.mark.parametrize(
+        ("statements", "expected"),
+        [
+            ([H(0)], [({0}, {0})]),
+            ([CNOT(0, 1)], [({0}, {0, 1})]),
+            ([H(0), CNOT(0, 1), H(1)], [({0, 1, 2}, {0, 1})]),
+            ([H(1), CNOT(0, 1), H(1), H(2), CNOT(1, 2), H(2)], [({0, 1, 2}, {0, 1}), ({3, 4, 5}, {1, 2})]),
+            (
+                [H(3), CNOT(2, 3), H(1), H(2), CNOT(0, 1), H(2), H(1), H(3), CNOT(1, 2), H(2)],
+                [({0, 1, 3, 5, 7}, {2, 3}), ({2, 4, 6}, {0, 1}), ({8, 9}, {1, 2})],
+            ),
+            (
+                [H(2), H(0), CNOT(0, 1), CNOT(1, 2), CNOT(2, 3), H(1), H(3), H(0), H(2)],
+                [({1, 2, 7}, {0, 1}), ({0, 3, 5}, {1, 2}), ({4, 6, 8}, {2, 3})],
+            ),
+        ],
+    )
+    def test_group_gates(self, statements: list[Statement], expected: list[tuple[set[int], set[int]]]) -> None:
+        ir = IR()
+        for statement in statements:
+            ir.add_statement(statement)
+
+        graph = build_graph(ir)
+
+        assert group_gates(graph) == expected
+
+
+@pytest.fixture
+def merger() -> TwoQubitGatesMerger:
+    return TwoQubitGatesMerger()
+
+
+@pytest.mark.parametrize(
+    ("circuit", "expected_circuit"),
+    [
+        (CircuitBuilder(1).H(0).to_circuit(), CircuitBuilder(1).H(0).to_circuit()),
+        (CircuitBuilder(2).CNOT(0, 1).to_circuit(), CircuitBuilder(2).CNOT(0, 1).to_circuit()),
+        (CircuitBuilder(2).CNOT(0, 1).CNOT(1, 0).CNOT(0, 1).to_circuit(), CircuitBuilder(2).SWAP(0, 1).to_circuit()),
+        (CircuitBuilder(2).SWAP(0, 1).SWAP(0, 1).to_circuit(), CircuitBuilder(2).I(0).to_circuit()),
+        (CircuitBuilder(2).SWAP(0, 1).CZ(0, 1).S(0).S(1).to_circuit(), CircuitBuilder(2).ISWAP(0, 1).to_circuit()),
+        (CircuitBuilder(2).H(1).CNOT(0, 1).H(1).to_circuit(), CircuitBuilder(2).CZ(0, 1).to_circuit()),
+        (CircuitBuilder(2).CV(0, 1).CV(0, 1).to_circuit(), CircuitBuilder(2).CNOT(0, 1).to_circuit()),
+        (CircuitBuilder(2).H(0).H(1).CNOT(0, 1).H(0).H(1).to_circuit(), CircuitBuilder(2).CNOT(1, 0).to_circuit()),
+        (
+            CircuitBuilder(2).T(0).H(1).CNOT(1, 0).Tdag(0).T(1).CNOT(1, 0).H(1).to_circuit(),
+            CircuitBuilder(2).CV(0, 1).to_circuit(),
+        ),
+    ],
+)
+def test_two_qubit_gates_merger_with_two_qubits(
+    circuit: Circuit, expected_circuit: Circuit, merger: TwoQubitGatesMerger
+) -> None:
+    expected_matrix = circuit_matrix_calculator.get_circuit_matrix(expected_circuit)
+    pre_merge_matrix = circuit_matrix_calculator.get_circuit_matrix(circuit)
+
+    merger.merge(circuit.ir, circuit.qubit_register_size)
+
+    actual_matrix = circuit_matrix_calculator.get_circuit_matrix(circuit)
+
+    assert are_matrices_equivalent_up_to_global_phase(actual_matrix, pre_merge_matrix)
+    assert are_matrices_equivalent_up_to_global_phase(actual_matrix, expected_matrix)
+
+    # Since we are only dealing two qubits, we can check that the number of statements is 1,
+    # which means that the two-qubit gates have been merged into a single gate.
+    assert len(circuit.ir.statements) == 1
+
+
+@pytest.mark.parametrize(
+    ("circuit", "expected_circuit"),
+    [
+        (
+            CircuitBuilder(3).H(1).CNOT(0, 1).H(1).H(2).CNOT(1, 2).H(2).to_circuit(),
+            CircuitBuilder(3).CZ(0, 1).CZ(1, 2).to_circuit(),
+        ),
+    ],
+)
+def test_two_qubit_gates_merger_with_multiple_qubits(
+    circuit: Circuit, expected_circuit: Circuit, merger: TwoQubitGatesMerger
+) -> None:
+    expected_matrix = circuit_matrix_calculator.get_circuit_matrix(expected_circuit)
+    pre_merge_matrix = circuit_matrix_calculator.get_circuit_matrix(circuit)
+
+    merger.merge(circuit.ir, circuit.qubit_register_size)
+
+    actual_matrix = circuit_matrix_calculator.get_circuit_matrix(circuit)
+
+    assert are_matrices_equivalent_up_to_global_phase(actual_matrix, pre_merge_matrix)
+    assert are_matrices_equivalent_up_to_global_phase(actual_matrix, expected_matrix)