Merge pull request #367 from qiboteam/pulsecompiler

Gate to pulse compiler

src/qibolab/backends.py

@@ -8,6 +8,7 @@
 from qibo.states import CircuitResult
 
 from qibolab import __version__ as qibolab_version
+from qibolab.compilers import Compiler
 from qibolab.platform import Platform
 from qibolab.platforms.abstract import AbstractPlatform
 from qibolab.transpilers import Pipeline
@@ -26,6 +27,7 @@ def __init__(self, platform, runcard=None):
             "numpy": self.np.__version__,
             "qibolab": qibolab_version,
         }
+        self.compiler = Compiler.default()
         self.transpiler = Pipeline.default(self.platform.two_qubit_natives)
 
     def apply_gate(self, gate, state, nqubits):  # pragma: no cover
@@ -34,6 +36,25 @@ def apply_gate(self, gate, state, nqubits):  # pragma: no cover
     def apply_gate_density_matrix(self, gate, state, nqubits):  # pragma: no cover
         raise_error(NotImplementedError, "Qibolab cannot apply gates directly.")
 
+    def assign_measurements(self, measurement_map, circuit_result):
+        """Assigning measurement outcomes to :class:`qibo.states.MeasurementResult` for each gate.
+
+        This allows properly obtaining the measured shots from the :class:`qibo.states.CircuitResult`
+        object returned by the circuit execution.
+
+        Args:
+            measurement_map (dict): Map from each measurement gate to the sequence of
+                readout pulses implementing it.
+            circuit_result (:class:`qibo.states.CircuitResult`): Circuit result object
+                containing the readout measurement shots. This is created in ``execute_circuit``.
+        """
+        readout = circuit_result.execution_result
+        for gate, sequence in measurement_map.items():
+            _samples = (readout[pulse.serial].shots for pulse in sequence.pulses)
+            samples = list(filter(lambda x: x is not None, _samples))
+            gate.result.backend = self
+            gate.result.register_samples(np.array(samples).T)
+
     def execute_circuit(
         self, circuit, initial_state=None, nshots=None, fuse_one_qubit=False, check_transpiled=False
     ):  # pragma: no cover
@@ -55,7 +76,7 @@ def execute_circuit(
             CircuitResult object containing the results acquired from the execution.
         """
         if isinstance(initial_state, type(circuit)):
-            self.execute_circuit(
+            return self.execute_circuit(
                 circuit=initial_state + circuit,
                 nshots=nshots,
                 fuse_one_qubit=fuse_one_qubit,
@@ -77,7 +98,7 @@ def execute_circuit(
                 self.transpiler.check_execution(circuit, native_circuit)
 
         # Transpile the native circuit into a sequence of pulses ``PulseSequence``
-        sequence = self.platform.transpile(native_circuit)
+        sequence, measurement_map = self.compiler.compile(native_circuit, self.platform)
 
         if not self.platform.is_connected:
             self.platform.connect()
@@ -87,19 +108,8 @@ def execute_circuit(
         self.platform.start()
         readout = self.platform.execute_pulse_sequence(sequence, nshots)
         self.platform.stop()
-        result = CircuitResult(self, native_circuit, readout, nshots)
-
-        # Register measurement outcomes
-        if isinstance(readout, dict):
-            for gate in native_circuit.queue:
-                if isinstance(gate, gates.M):
-                    samples = []
-                    for serial in gate.pulses:
-                        shots = readout[serial].shots
-                        if shots is not None:
-                            samples.append(shots)
-                    gate.result.backend = self
-                    gate.result.register_samples(np.array(samples).T)
+        result = CircuitResult(self, circuit, readout, nshots)
+        self.assign_measurements(measurement_map, result)
         return result
 
     def circuit_result_tensor(self, result):

src/qibolab/compilers/__init__.py

@@ -0,0 +1 @@
+from qibolab.compilers.compiler import Compiler

src/qibolab/compilers/compiler.py

@@ -0,0 +1,146 @@
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+from qibo import gates
+from qibo.config import raise_error
+
+from qibolab.compilers.default import (
+    cz_rule,
+    identity_rule,
+    measurement_rule,
+    rz_rule,
+    u3_rule,
+    z_rule,
+)
+from qibolab.pulses import PulseSequence, ReadoutPulse
+
+
+@dataclass
+class Compiler:
+    """Compiler that transforms a :class:`qibo.models.Circuit` to a :class:`qibolab.pulses.PulseSequence`.
+
+    The transformation is done using a dictionary of rules which map each Qibo gate to a
+    pulse sequence and some virtual Z-phases.
+
+    A rule is a function that takes two argumens:
+        gate (:class:`qibo.gates.abstract.Gate`): Gate object to be compiled.
+        platform (:class:`qibolab.platforms.abstract.AbstractPlatform`): Platform object to read
+            native gate pulses from.
+    and returns
+        sequence (:class:`qibolab.pulses.PulseSequence`): Sequence of pulses that implement
+            the given gate.
+        virtual_z_phases (dict): Dictionary mapping qubits to virtual Z-phases induced by the gate.
+
+    See ``qibolab.compilers.default`` for an example of a compiler implementation.
+    """
+
+    rules: dict = field(default_factory=dict)
+    """Map from gates to compilation rules."""
+
+    @classmethod
+    def default(cls):
+        return cls(
+            {
+                gates.I: identity_rule,
+                gates.Z: z_rule,
+                gates.RZ: rz_rule,
+                gates.U3: u3_rule,
+                gates.CZ: cz_rule,
+                gates.M: measurement_rule,
+            }
+        )
+
+    def __setitem__(self, key, rule):
+        """Sets a new rule to the compiler.
+
+        If a rule already exists for the gate, it will be overwritten.
+        """
+        self.rules[key] = rule
+
+    def __getitem__(self, item):
+        """Get an existing rule for a given gate."""
+        try:
+            return self.rules[item]
+        except KeyError:
+            raise_error(KeyError, f"Compiler rule not available for {item}.")
+
+    def __delitem__(self, item):
+        """Remove rule for the given gate."""
+        try:
+            del self.rules[item]
+        except KeyError:
+            raise_error(KeyError, f"Cannot remove {item} from compiler because it does not exist.")
+
+    def register(self, gate_cls):
+        """Decorator for registering a function as a rule in the compiler.
+
+        Using this decorator is optional. Alternatively the user can set the rules directly
+        via ``__setitem__`.
+
+        Args:
+            gate_cls: Qibo gate object that the rule will be assigned to.
+        """
+
+        def inner(func):
+            self[gate_cls] = func
+            return func
+
+        return inner
+
+    def _compile_gate(self, gate, platform, sequence, virtual_z_phases, moment_start):
+        """Adds a single gate to the pulse sequence."""
+        rule = self[gate.__class__]
+        # get local sequence and phases for the current gate
+        gate_sequence, gate_phases = rule(gate, platform)
+
+        # update global pulse sequence
+        # determine the right start time based on the availability of the qubits involved
+        all_qubits = {*gate_sequence.qubits, *gate.qubits}
+        start = max(sequence.get_qubit_pulses(*all_qubits).finish, moment_start)
+        # shift start time and phase according to the global sequence
+        for pulse in gate_sequence:
+            pulse.start += start
+            if not isinstance(pulse, ReadoutPulse):
+                pulse.relative_phase += virtual_z_phases[pulse.qubit]
+            sequence.add(pulse)
+
+        return gate_sequence, gate_phases
+
+    def compile(self, circuit, platform):
+        """Transforms a circuit to pulse sequence.
+
+        Args:
+            circuit (qibo.models.Circuit): Qibo circuit that respects the platform's
+                connectivity and native gates.
+            platform (qibolab.platforms.abstract.AbstractPlatform): Platform used
+                to load the native pulse representations.
+
+        Returns:
+            sequence (qibolab.pulses.PulseSequence): Pulse sequence that implements the circuit.
+            measurement_map (dict): Map from each measurement gate to the sequence of  readout pulses
+                implementing it.
+        """
+        sequence = PulseSequence()
+        # FIXME: This will not work with qubits that have string names
+        # TODO: Implement a mapping between circuit qubit ids and platform ``Qubit``s
+        virtual_z_phases = defaultdict(int)
+
+        measurement_map = {}
+        # process circuit gates
+        for moment in circuit.queue.moments:
+            moment_start = sequence.finish
+            for gate in set(filter(lambda x: x is not None, moment)):
+                gate_sequence, gate_phases = self._compile_gate(
+                    gate, platform, sequence, virtual_z_phases, moment_start
+                )
+
+                # update virtual Z phases
+                for qubit, phase in gate_phases.items():
+                    virtual_z_phases[qubit] += phase
+
+                # register readout sequences to ``measurement_map`` so that we can
+                # properly map acquisition results to measurement gates
+                if isinstance(gate, gates.M):
+                    measurement_map[gate] = gate_sequence
+
+        return sequence, measurement_map

src/qibolab/compilers/default.py

@@ -0,0 +1,68 @@
+"""Implementation of the default compiler.
+
+Uses I, Z, RZ, U3, CZ and M as the set of native gates.
+"""
+import math
+
+from qibolab.pulses import PulseSequence
+
+
+def identity_rule(gate, platform):
+    """Identity gate skipped."""
+    return PulseSequence(), {}
+
+
+def z_rule(gate, platform):
+    """Z gate applied virtually."""
+    qubit = gate.target_qubits[0]
+    return PulseSequence(), {qubit: math.pi}
+
+
+def rz_rule(gate, platform):
+    """RZ gate applied virtually."""
+    qubit = gate.target_qubits[0]
+    return PulseSequence(), {qubit: gate.parameters[0]}
+
+
+def u3_rule(gate, platform):
+    """U3 applied as RZ-RX90-RZ-RX90-RZ."""
+    qubit = gate.target_qubits[0]
+    # Transform gate to U3 and add pi/2-pulses
+    theta, phi, lam = gate.parameters
+    # apply RZ(lam)
+    virtual_z_phases = {qubit: lam}
+    sequence = PulseSequence()
+    # Fetch pi/2 pulse from calibration
+    RX90_pulse_1 = platform.create_RX90_pulse(qubit, start=0, relative_phase=virtual_z_phases[qubit])
+    # apply RX(pi/2)
+    sequence.add(RX90_pulse_1)
+    # apply RZ(theta)
+    virtual_z_phases[qubit] += theta
+    # Fetch pi/2 pulse from calibration
+    RX90_pulse_2 = platform.create_RX90_pulse(
+        qubit, start=RX90_pulse_1.finish, relative_phase=virtual_z_phases[qubit] - math.pi
+    )
+    # apply RX(-pi/2)
+    sequence.add(RX90_pulse_2)
+    # apply RZ(phi)
+    virtual_z_phases[qubit] += phi
+
+    return sequence, virtual_z_phases
+
+
+def cz_rule(gate, platform):
+    """CZ applied as defined in the platform runcard.
+
+    Applying the CZ gate may involve sending pulses on qubits
+    that the gate is not directly acting on.
+    """
+    return platform.create_CZ_pulse_sequence(gate.qubits)
+
+
+def measurement_rule(gate, platform):
+    """Measurement gate applied using the platform readout pulse."""
+    sequence = PulseSequence()
+    for qubit in gate.target_qubits:
+        MZ_pulse = platform.create_MZ_pulse(qubit, start=0)
+        sequence.add(MZ_pulse)
+    return sequence, {}