Create merge pass

- Circuit.merge_single_qubit_gates
- mcKay decomposer is now only the decomposition
- Circuit.decompose is introduced, pass it the McKay decomposer

Author: Pablo Le Hénaff

opensquirrel/circuit.py

@@ -1,9 +1,9 @@
-from typing import Callable, Dict
+from typing import Callable, Dict, List
 
 import numpy as np
 
 import opensquirrel.parsing.antlr.squirrel_ir_from_string
-from opensquirrel import circuit_matrix_calculator, mckay_decomposer, replacer, writer
+from opensquirrel import circuit_matrix_calculator, replacer, mckay_decomposer, merger, writer
 from opensquirrel.default_gates import default_gate_aliases, default_gate_set
 from opensquirrel.parsing.libqasm.libqasm_ir_creator import LibqasmIRCreator
 from opensquirrel.squirrel_ir import Gate, SquirrelIR
@@ -21,7 +21,7 @@ class Circuit:
         <BLANKLINE>
         h q[0]
         <BLANKLINE>
-        >>> c.decompose_mckay()
+        >>> c.decompose(decomposer=mckay_decomposer.mckay_decomposer)
         >>> c
         version 3.0
         <BLANKLINE>
@@ -84,27 +84,25 @@ def number_of_qubits(self) -> int:
     def qubit_register_name(self) -> str:
         return self.squirrel_ir.qubit_register_name
 
-    def decompose_mckay(self):
-        """Perform gate fusion on all one-qubit gates and decompose them in the McKay style.
-
-        * all one-qubit gates on same qubit are merged together, without attempting to commute any gate
-        * two-or-more-qubit gates are left as-is
-        * merged one-qubit gates are decomposed according to McKay decomposition, that is:
-                gate   ---->    Rz.Rx(pi/2).Rz.Rx(pi/2).Rz
-        * _global phase is deemed irrelevant_, therefore a simulator backend might produce different output
-            for the input and output circuit - those outputs should be equivalent modulo global phase.
+    def merge_single_qubit_gates(self):
+        """Perform gate fusion on all one-qubit gates.
+        Gates obtained from merging other gates become anonymous gates.
         """
 
-        self.squirrel_ir = mckay_decomposer.decompose_mckay(self.squirrel_ir)  # FIXME: inplace
+        merger.merge_single_qubit_gates(self.squirrel_ir)
 
-    def replace(self, gate_name: str, f):
-        """Manually replace occurrences of a given gate with a list of gates.
+    def decompose(self, decomposer: Callable[[Gate], List[Gate]]):
+        """Generic decomposition pass. It applies the given decomposer function
+        to every gate in the circuit."""
+        replacer.decompose(self.squirrel_ir, decomposer)
 
-        * this can be called decomposition - but it's the least fancy version of it
-        * function parameter gives the decomposition based on parameters of original gate
+    def replace(self, gate_generator: Callable[..., Gate], f):
+        """Manually replace occurrences of a given gate with a list of gates.
+        `f` is a callable that takes the arguments of the replaced gate
+        and returns the decomposition as a list of gates.
         """
 
-        replacer.replace(self.squirrel_ir, gate_name, f)
+        replacer.replace(self.squirrel_ir, gate_generator, f)
 
     def test_get_circuit_matrix(self) -> np.ndarray:
         """Get the (large) unitary matrix corresponding to the circuit.
@@ -117,9 +115,6 @@ def test_get_circuit_matrix(self) -> np.ndarray:
         return circuit_matrix_calculator.get_circuit_matrix(self.squirrel_ir)
 
     def __repr__(self) -> str:
-        """Write the circuit to a cQasm3 string.
-
-        * comments are removed
-        """
+        """Write the circuit to a cQasm3 string."""
 
         return writer.squirrel_ir_to_string(self.squirrel_ir)

opensquirrel/default_gates.py

@@ -18,6 +18,11 @@ def x90(q: Qubit) -> Gate:
     return BlochSphereRotation(qubit=q, axis=(1, 0, 0), angle=math.pi / 2, phase=0)
 
 
+@named_gate
+def xm90(q: Qubit) -> Gate:
+    return BlochSphereRotation(qubit=q, axis=(1, 0, 0), angle=-math.pi / 2, phase=0)
+
+
 @named_gate
 def y(q: Qubit) -> Gate:
     return BlochSphereRotation(qubit=q, axis=(0, 1, 0), angle=math.pi, phase=math.pi / 2)
@@ -76,5 +81,5 @@ def crk(control: Qubit, target: Qubit, k: Int) -> Gate:
     return ControlledGate(control, BlochSphereRotation(qubit=target, axis=(0, 0, 1), angle=theta, phase=theta / 2))
 
 
-default_gate_set = [h, x, x90, y, y90, z, z90, cz, cr, crk, cnot, rx, ry, rz, x]
+default_gate_set = [h, x, x90, xm90, y, y90, z, z90, cz, cr, crk, cnot, rx, ry, rz, x]
 default_gate_aliases = {"X": x, "RX": rx, "RY": ry, "RZ": rz, "Hadamard": h, "H": h}

opensquirrel/mckay_decomposer.py

@@ -8,110 +8,49 @@
 from opensquirrel.squirrel_ir import BlochSphereRotation, Float, Gate, Qubit, SquirrelIR
 
 
-class _McKayDecomposerImpl:
-    def __init__(self, number_of_qubits: int, qubit_register_name: str):
-        self.output = SquirrelIR(number_of_qubits=number_of_qubits, qubit_register_name=qubit_register_name)
-        self.accumulated_1q_gates = {}
+def mckay_decomposer(g: Gate):
+    """Return the McKay decomposition of a 1-qubit gate as a list of gates.
+            gate   ---->    Rz.Rx(pi/2).Rz.Rx(pi/2).Rz
 
-    def decompose_and_add(self, qubit: Qubit, angle: float, axis: Tuple[float, float, float]):
-        if abs(angle) < ATOL:
-            return
+    _global phase is deemed irrelevant_, therefore a simulator backend might produce different output
+        for the input and output - the results should be equivalent modulo global phase.
+    """
+    if not isinstance(g, BlochSphereRotation):
+        return [g]
 
-        # McKay decomposition
+    if abs(g.angle) < ATOL:
+        return []
 
-        za_mod = sqrt(cos(angle / 2) ** 2 + (axis[2] * sin(angle / 2)) ** 2)
-        zb_mod = abs(sin(angle / 2)) * sqrt(axis[0] ** 2 + axis[1] ** 2)
+    # McKay decomposition
 
-        theta = pi - 2 * atan2(zb_mod, za_mod)
+    za_mod = sqrt(cos(g.angle / 2) ** 2 + (g.axis[2] * sin(g.angle / 2)) ** 2)
+    zb_mod = abs(sin(g.angle / 2)) * sqrt(g.axis[0] ** 2 + g.axis[1] ** 2)
 
-        alpha = atan2(-sin(angle / 2) * axis[2], cos(angle / 2))
-        beta = atan2(-sin(angle / 2) * axis[0], -sin(angle / 2) * axis[1])
+    theta = pi - 2 * atan2(zb_mod, za_mod)
 
-        lam = beta - alpha
-        phi = -beta - alpha - pi
+    alpha = atan2(-sin(g.angle / 2) * g.axis[2], cos(g.angle / 2))
+    beta = atan2(-sin(g.angle / 2) * g.axis[0], -sin(g.angle / 2) * g.axis[1])
 
-        lam = normalize_angle(lam)
-        phi = normalize_angle(phi)
-        theta = normalize_angle(theta)
+    lam = beta - alpha
+    phi = -beta - alpha - pi
 
-        if abs(lam) > ATOL:
-            self.output.add_gate(rz(qubit, Float(lam)))
+    lam = normalize_angle(lam)
+    phi = normalize_angle(phi)
+    theta = normalize_angle(theta)
 
-        self.output.add_gate(x90(qubit))
+    decomposed_g = []
 
-        if abs(theta) > ATOL:
-            self.output.add_gate(rz(qubit, Float(theta)))
+    if abs(lam) > ATOL:
+        decomposed_g.append(rz(g.qubit, Float(lam)))
 
-        self.output.add_gate(x90(qubit))
+    decomposed_g.append(x90(g.qubit))
 
-        if abs(phi) > ATOL:
-            self.output.add_gate(rz(qubit, Float(phi)))
+    if abs(theta) > ATOL:
+        decomposed_g.append(rz(g.qubit, Float(theta)))
 
-    def flush(self, q):
-        if q not in self.accumulated_1q_gates:
-            return
-        p = self.accumulated_1q_gates.pop(q)
-        self.decompose_and_add(q, p["angle"], p["axis"])
+    decomposed_g.append(x90(g.qubit))
 
-    def flush_all(self):
-        while len(self.accumulated_1q_gates) > 0:
-            self.flush(next(iter(self.accumulated_1q_gates)))
+    if abs(phi) > ATOL:
+        decomposed_g.append(rz(g.qubit, Float(phi)))
 
-    def accumulate(self, qubit, bloch_sphere_rotation: BlochSphereRotation):
-        axis, angle, phase = bloch_sphere_rotation.axis, bloch_sphere_rotation.angle, bloch_sphere_rotation.phase
-
-        if qubit not in self.accumulated_1q_gates:
-            self.accumulated_1q_gates[qubit] = {"angle": angle, "axis": axis, "phase": phase}
-            return
-
-        existing = self.accumulated_1q_gates[qubit]
-        combined_phase = phase + existing["phase"]
-
-        a = angle
-        l = axis
-        b = existing["angle"]
-        m = existing["axis"]
-
-        combined_angle = 2 * acos(cos(a / 2) * cos(b / 2) - sin(a / 2) * sin(b / 2) * np.dot(l, m))
-
-        if abs(sin(combined_angle / 2)) < ATOL:
-            self.accumulated_1q_gates.pop(qubit)
-            return
-
-        combined_axis = (
-            1
-            / sin(combined_angle / 2)
-            * (sin(a / 2) * cos(b / 2) * l + cos(a / 2) * sin(b / 2) * m + sin(a / 2) * sin(b / 2) * np.cross(l, m))
-        )
-
-        self.accumulated_1q_gates[qubit] = {"angle": combined_angle, "axis": combined_axis, "phase": combined_phase}
-
-    def process_gate(self, gate: Gate):
-        qubit_arguments = [arg for arg in gate.arguments if isinstance(arg, Qubit)]
-
-        if len(qubit_arguments) >= 2:
-            [self.flush(q) for q in qubit_arguments]
-            self.output.add_gate(gate)
-            return
-
-        if len(qubit_arguments) == 0:
-            assert False, "Unsupported"
-            return
-
-        assert isinstance(gate, BlochSphereRotation), f"Not supported for single qubit gate `{gate.name}`: {type(gate)}"
-
-        self.accumulate(qubit_arguments[0], gate)
-
-
-def decompose_mckay(squirrel_ir):
-    impl = _McKayDecomposerImpl(squirrel_ir.number_of_qubits, squirrel_ir.qubit_register_name)
-
-    for statement in squirrel_ir.statements:
-        if not isinstance(statement, Gate):
-            continue
-
-        impl.process_gate(statement)
-
-    impl.flush_all()
-
-    return impl.output  # FIXME: instead of returning a new IR, modify existing one
+    return decomposed_g

opensquirrel/merger.py

@@ -0,0 +1,84 @@
+from math import acos, atan2, cos, pi, sin, sqrt
+from typing import List, Optional
+
+import numpy as np
+
+from opensquirrel.common import ATOL
+from opensquirrel.squirrel_ir import BlochSphereRotation, Gate, Qubit, SquirrelIR
+
+
+def compose_bloch_sphere_rotations(a: BlochSphereRotation, b: BlochSphereRotation) -> BlochSphereRotation:
+    """Computes the Bloch sphere rotation resulting from the composition of two Block sphere rotations.
+    The first rotation is applied and then the second.
+    The resulting gate is anonymous except if `a` is the identity and `b` is not anonymous, or vice versa.
+
+    Uses Rodrigues' rotation formula, see for instance https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula.
+    """
+    assert a.qubit == b.qubit, "Cannot merge two BlochSphereRotation's on different qubits"
+
+    combined_angle = 2 * acos(
+        cos(a.angle / 2) * cos(b.angle / 2) - sin(a.angle / 2) * sin(b.angle / 2) * np.dot(a.axis, b.axis)
+    )
+
+    if abs(sin(combined_angle / 2)) < ATOL:
+        return BlochSphereRotation.identity(a.qubit)
+
+    combined_axis = (
+        1
+        / sin(combined_angle / 2)
+        * (
+            sin(a.angle / 2) * cos(b.angle / 2) * a.axis
+            + cos(a.angle / 2) * sin(b.angle / 2) * b.axis
+            + sin(a.angle / 2) * sin(b.angle / 2) * np.cross(a.axis, b.axis)
+        )
+    )
+
+    combined_phase = a.phase + b.phase
+
+    generator = b.generator if a.is_identity() else a.generator if b.is_identity() else None
+    arguments = b.arguments if a.is_identity() else a.arguments if b.is_identity() else None
+
+    return BlochSphereRotation(
+        qubit=a.qubit,
+        axis=combined_axis,
+        angle=combined_angle,
+        phase=combined_phase,
+        generator=generator,
+        arguments=arguments,
+    )
+
+
+def merge_single_qubit_gates(squirrel_ir: SquirrelIR):
+    accumulators_per_qubit: dict[Qubit, BlochSphereRotation] = {
+        Qubit(q): BlochSphereRotation.identity(Qubit(q)) for q in range(squirrel_ir.number_of_qubits)
+    }
+
+    statement_index = 0
+    while statement_index < len(squirrel_ir.statements):
+        statement = squirrel_ir.statements[statement_index]
+
+        if not isinstance(statement, Gate):
+            # Skip
+            statement_index += 1
+            continue
+
+        if isinstance(statement, BlochSphereRotation):
+            # Accumulate
+            already_accumulated = accumulators_per_qubit.get(statement.qubit)
+
+            composed = compose_bloch_sphere_rotations(statement, already_accumulated)
+            accumulators_per_qubit[statement.qubit] = composed
+
+            del squirrel_ir.statements[statement_index]
+            continue
+
+        for qubit_operand in statement.get_qubit_operands():
+            if not accumulators_per_qubit[qubit_operand].is_identity():
+                squirrel_ir.statements.insert(statement_index, accumulators_per_qubit[qubit_operand])
+                accumulators_per_qubit[qubit_operand] = BlochSphereRotation.identity(qubit_operand)
+                statement_index += 1
+        statement_index += 1
+
+    for accumulated_bloch_sphere_rotation in accumulators_per_qubit.values():
+        if not accumulated_bloch_sphere_rotation.is_identity():
+            squirrel_ir.statements.append(accumulated_bloch_sphere_rotation)

opensquirrel/replacer.py

@@ -1,31 +1,34 @@
-from typing import List
+from typing import Callable, List
 
-from opensquirrel.squirrel_ir import Comment, Gate, SquirrelIR
+from opensquirrel.squirrel_ir import Gate, SquirrelIR
 
 
-def replace(squirrel_ir: SquirrelIR, gate_name_to_replace: str, f):
+def decompose(squirrel_ir: SquirrelIR, decomposer: Callable[[Gate], List[Gate]]):
+    """Applies `decomposer` to every gate in the circuit, replacing each gate by the output of `decomposer`.
+    When `decomposer` decides to not decompose a gate, it needs to return a list with the intact gate as single element.
+    """
     statement_index = 0
     while statement_index < len(squirrel_ir.statements):
         statement = squirrel_ir.statements[statement_index]
 
-        if isinstance(statement, Comment):
-            statement_index += 1
-            continue
-
         if not isinstance(statement, Gate):
-            raise Exception("Unsupported")
-
-        if statement.name != gate_name_to_replace:
             statement_index += 1
             continue
 
-        # FIXME: handle case where if f is not a function but directly a list.
-
-        replacement: List[Gate] = f(*statement.arguments)
+        replacement: List[Gate] = decomposer(statement)
         squirrel_ir.statements[statement_index : statement_index + 1] = replacement
         statement_index += len(replacement)
 
         # TODO: Here, check that the semantic of the replacement is the same!
         # For this, need to update the simulation capabilities.
 
-        # TODO: Do we allow skipping the replacement, based on arguments?
+
+def replace(squirrel_ir: SquirrelIR, gate_generator: Callable[..., Gate], f):
+    """Does the same as decomposer, but only applies to a given gate."""
+
+    def generic_replacer(g: Gate) -> [Gate]:
+        if g.is_anonymous or g.generator != gate_generator:
+            return [g]
+        return f(*g.arguments)
+
+    decompose(squirrel_ir, generic_replacer)