feat: Allow passing PassManager (#231)

Author: speller26

qiskit_braket_provider/providers/adapter.py

@@ -12,7 +12,7 @@
 from qiskit.circuit import Instruction as QiskitInstruction
 from qiskit.circuit.library import get_standard_gate_name_mapping
 from qiskit.circuit.parametervector import ParameterVectorElement
-from qiskit.transpiler import Target
+from qiskit.transpiler import PassManager, Target
 from qiskit_ionq import add_equivalences, ionq_gates
 from sympy import Add, Mul, Pow, Symbol
 
@@ -499,6 +499,8 @@ def to_braket(
     target: Target | None = None,
     qubit_labels: Sequence[int] | None = None,
     optimization_level: int = 0,
+    callback: Callable | None = None,
+    pass_manager: PassManager | None = None,
 ) -> Circuit:
     """Return a Braket quantum circuit from a Qiskit quantum circuit.
 
@@ -522,6 +524,10 @@ def to_braket(
             assumed to be contiguous.
         optimization_level (int): The optimization level to pass to `qiskit.transpile`.
             Default: 0 (no optimization).
+        callback (Callable | None): A callback function that will be called after each transpiler
+            pass execution. Default: `None`.
+        pass_manager (PassManager): `PassManager` to transpile the circuit; will raise an error if
+            used in conjunction with a target, basis gates, or connectivity. Default: `None`.
 
     Returns:
         Circuit: Braket circuit
@@ -530,30 +536,41 @@ def to_braket(
         circuit = to_qiskit(circuit)
     if not isinstance(circuit, QuantumCircuit):
         raise TypeError(f"Expected a QuantumCircuit, got {type(circuit)} instead.")
-    if (basis_gates or connectivity) and target:
-        raise ValueError("Basis gates and connectivity cannot be specified alongside target.")
-
-    # If basis_gates is not None, then target remains empty
-    target = target if basis_gates or target else _create_default_target(circuit)
-    needs_transpilation = (
-        target
-        or connectivity
-        or (basis_gates and not {gate.name for gate, _, _ in circuit.data}.issubset(basis_gates))
-    )
-    if not verbatim and needs_transpilation:
-        circuit = transpile(
-            circuit,
-            basis_gates=basis_gates,
-            coupling_map=connectivity,
-            optimization_level=optimization_level,
-            target=target,
+    loose_constraints = basis_gates or connectivity
+    if pass_manager and (target or loose_constraints):
+        raise ValueError(
+            "Cannot specify target, basis gates, or connectivity alongside pass manager"
         )
+    if loose_constraints and target:
+        raise ValueError("Cannot specify basis gates or connectivity alongside target.")
+
+    if pass_manager:
+        circuit = pass_manager.run(circuit, callback=callback)
+    elif not verbatim:
+        # If basis_gates is not None, then target remains empty
+        target = target if basis_gates or target else _default_target(circuit)
+        if (
+            target
+            or connectivity
+            or (
+                basis_gates and not {gate.name for gate, _, _ in circuit.data}.issubset(basis_gates)
+            )
+        ):
+            circuit = transpile(
+                circuit,
+                basis_gates=basis_gates,
+                coupling_map=connectivity,
+                optimization_level=optimization_level,
+                target=target,
+                callback=callback,
+            )
+
     # Verify that ParameterVector would not collide with scalar variables after renaming.
     _validate_name_conflicts(circuit.parameters)
     # Handle qiskit to braket conversion
     measured_qubits: dict[int, int] = {}
     braket_circuit = Circuit()
-    qubit_labels = qubit_labels or sorted(circuit.find_bit(q).index for q in circuit.qubits)
+    qubit_labels = qubit_labels or _default_qubit_labels(circuit)
     for circuit_instruction in circuit.data:
         operation = circuit_instruction.operation
         qubits = circuit_instruction.qubits
@@ -630,7 +647,7 @@ def to_braket(
 
     # QPU targets will have qubits/pairs specified for each instruction;
     # Targets whose values consist solely of {None: None} are either simulator or default targets
-    if verbatim or (target and any(v != {None: None} for v in target.values())):
+    if verbatim or (target and any(v != {None: None} for v in target.values())) or pass_manager:
         braket_circuit = Circuit(braket_circuit.result_types).add_verbatim_box(
             Circuit(braket_circuit.instructions)
         )
@@ -641,15 +658,20 @@ def to_braket(
     return braket_circuit
 
 
-def _create_default_target(circuit: QuantumCircuit) -> Target:
+def _default_target(circuit: QuantumCircuit) -> Target:
     target = Target(num_qubits=circuit.num_qubits)
     for braket_name, instruction in _BRAKET_GATE_NAME_TO_QISKIT_GATE.items():
-        if (name := braket_name.lower()) in _BRAKET_TO_QISKIT_NAMES:
-            target.add_instruction(instruction, name=_BRAKET_TO_QISKIT_NAMES[name])
+        if name := _BRAKET_TO_QISKIT_NAMES.get(braket_name.lower()):
+            target.add_instruction(instruction, name=name)
     target.add_instruction(Measure())
     return target
 
 
+def _default_qubit_labels(circuit: QuantumCircuit) -> tuple[int, ...]:
+    bits = sorted(circuit.find_bit(q).index for q in circuit.qubits)
+    return tuple(range(max(bits) + 1)) if bits else tuple()
+
+
 def _create_free_parameters(operation):
     params = operation.params if hasattr(operation, "params") else []
     for i, param in enumerate(params):

qiskit_braket_provider/providers/braket_backend.py

@@ -5,11 +5,12 @@
 import logging
 import warnings
 from abc import ABC
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from typing import Generic, TypeVar
 
 from qiskit import QuantumCircuit
 from qiskit.providers import BackendV2, Options, QubitProperties
+from qiskit.transpiler import PassManager, Target
 
 from braket.aws import AwsDevice, AwsDeviceType, AwsQuantumTask
 from braket.aws.queue_information import QueueDepthInfo
@@ -343,6 +344,8 @@ def run(
         native: bool = False,
         *,
         optimization_level: int = 0,
+        callback: Callable | None = None,
+        pass_manager: PassManager | None = None,
         **options,
     ):
         if isinstance(run_input, QuantumCircuit):
@@ -357,13 +360,12 @@ def run(
             del options["meas_level"]
 
         # Always use target for simulator
-        target, basis_gates = (
-            (self._target, None)
-            if native or self._device.type == AwsDeviceType.SIMULATOR
-            else (None, self._gateset)
-        )
+        target, basis_gates = self._target_and_basis_gates(native, pass_manager)
         braket_circuits = (
-            [to_braket(circ, verbatim=True, qubit_labels=self._qubit_labels) for circ in circuits]
+            [
+                to_braket(circ, verbatim=True, qubit_labels=self._qubit_labels, callback=callback)
+                for circ in circuits
+            ]
             if verbatim
             else [
                 to_braket(
@@ -375,6 +377,8 @@ def run(
                         native_angle_restrictions(self._device.properties) if native else None
                     ),
                     optimization_level=optimization_level,
+                    callback=callback,
+                    pass_manager=pass_manager,
                 )
                 for circ in circuits
             ]
@@ -386,6 +390,16 @@ def run(
             else self._run_batch(braket_circuits, shots, **options)
         )
 
+    def _target_and_basis_gates(
+        self, native: bool, pass_manager: PassManager
+    ) -> tuple[Target | None, set[str] | None]:
+        if pass_manager:
+            return None, None
+        if native or self._device.type == AwsDeviceType.SIMULATOR:
+            # Always use target for simulator
+            return self._target, None
+        return None, self._gateset
+
     def _run_batch(self, braket_circuits: list[Circuit], shots: int, **options):
         batch_task = self._device.run_batch(braket_circuits, shots=shots, **options)
         tasks: list[AwsQuantumTask] = batch_task.tasks

tests/providers/test_adapter.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import pytest
-from qiskit import ClassicalRegister, QuantumCircuit, QuantumRegister
+from qiskit import ClassicalRegister, QuantumCircuit, QuantumRegister, generate_preset_pass_manager
 from qiskit.circuit import Instruction as QiskitInstruction
 from qiskit.circuit import Parameter, ParameterVector
 from qiskit.circuit.library import GlobalPhaseGate, PauliEvolutionGate
@@ -274,9 +274,28 @@ def test_target_with_loose_constraints(self):
         target.add_instruction(qiskit_gates.HGate())
 
         with pytest.raises(ValueError):
-            to_braket(circuit, basis_gates={"h"}, target=target)
+            to_braket(circuit, target=target, basis_gates={"h"})
         with pytest.raises(ValueError):
-            to_braket(circuit, connectivity=[[0, 1], [1, 2]], target=target)
+            to_braket(circuit, target=target, connectivity=[[0, 1], [1, 2]])
+
+    def test_pass_manager_with_other_arguments(self):
+        """
+        Tests that to_braket raises a ValueError if pass_manager is supplied
+        with target or loose constraints.
+        """
+        circuit = QuantumCircuit(1, 1)
+        circuit.h(0)
+
+        target = Target()
+        target.add_instruction(qiskit_gates.HGate())
+        pass_manager = generate_preset_pass_manager(2, target=target)
+
+        with pytest.raises(ValueError):
+            to_braket(circuit, pass_manager=pass_manager, target=target)
+        with pytest.raises(ValueError):
+            to_braket(circuit, pass_manager=pass_manager, basis_gates={"h"})
+        with pytest.raises(ValueError):
+            to_braket(circuit, pass_manager=pass_manager, connectivity=[[0, 1], [1, 2]])
 
     def test_convert_parametric_qiskit_to_braket_circuit(self):
         """Tests to_braket works with parametric circuits."""

tests/providers/test_braket_backend.py

@@ -7,7 +7,7 @@
 import numpy as np
 from botocore import errorfactory
 from networkx import DiGraph, complete_graph, from_dict_of_lists, relabel_nodes
-from qiskit import QuantumCircuit, transpile
+from qiskit import QuantumCircuit, generate_preset_pass_manager, transpile
 from qiskit.circuit import Instruction as QiskitInstruction
 from qiskit.circuit.library import TwoLocal
 from qiskit.circuit.random import random_circuit
@@ -258,6 +258,29 @@ def test_run_multiple_circuits_program_set(self):
         )
         self.assertEqual(device.run.call_count, 2)
 
+    def test_run_with_pass_manager(self):
+        """Tests run with pass_manager"""
+        device = Mock()
+        device.properties = RIGETTI_MOCK_GATE_MODEL_QPU_CAPABILITIES
+        device.type = "QPU"
+        device.topology_graph = None
+        backend = BraketAwsBackend(device=device)
+        mock_task_1 = Mock(spec=LocalQuantumTask)
+        mock_task_1.id = "0"
+        mock_task_2 = Mock(spec=LocalQuantumTask)
+        mock_task_2.id = "1"
+        mock_batch = Mock(spec=AwsQuantumTaskBatch)
+        mock_batch.tasks = [mock_task_1, mock_task_2]
+        backend._device.run_batch = Mock(return_value=mock_batch)
+        circuit = QuantumCircuit(1)
+        circuit.h(0)
+
+        backend.run(circuit, shots=0, pass_manager=generate_preset_pass_manager(2, backend))
+        native_circuit = Circuit().add_verbatim_box(
+            Circuit().rz(0, np.pi / 2).rx(0, np.pi / 2).rz(0, np.pi / 2)
+        )
+        device.run_batch.assert_called_once_with([native_circuit], shots=0)
+
     def test_run_invalid_run_input(self):
         """Tests run with invalid input to run"""
         device = Mock()

tests/providers/test_braket_instructions.py

@@ -7,7 +7,7 @@
 
 from braket.experimental_capabilities import EnableExperimentalCapability
 from qiskit_braket_provider import to_braket
-from qiskit_braket_provider.providers.adapter import _create_default_target
+from qiskit_braket_provider.providers.adapter import _default_target
 from qiskit_braket_provider.providers.braket_instructions import CCPRx, MeasureFF
 
 
@@ -83,7 +83,7 @@ def test_circuit_with_measureff_ccprx(self):
             CCPRx(0.5, 0.7, 0), qubits=(Qubit(QuantumRegister(1, "q"), 0),)
         )
 
-        target = _create_default_target(circuit)
+        target = _default_target(circuit)
         target.add_instruction(
             CCPRx(Parameter("angle_1"), Parameter("angle_2"), Parameter("feedback_key"))
         )