feature: device-specific gate validation for autoqasm programs

.github/workflows/check-format.yml

@@ -28,4 +28,4 @@ jobs:
     - name: Run code format checks
       run: |
         black --check . --extend-exclude=src/braket/experimental/autoqasm/autograph
-        flake8 --enable-extensions=BCS src/braket/experimental/autoqasm
+        flake8 --enable-extensions=BCS src

examples/autoqasm/ionq_gates.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+import braket.experimental.autoqasm as aq
+from braket.experimental.autoqasm.instructions import gpi, gpi2, ms
+
+
+@aq.gate
+def h(q: aq.Qubit):
+    gpi2(q, np.pi / 2)
+    gpi(q, 0)
+
+
+@aq.gate
+def u(q: aq.Qubit, a: float, b: float, c: float):
+    gpi2(q, a)
+    gpi(q, b)
+    gpi2(q, c)
+
+
+@aq.gate
+def rx(q: aq.Qubit, theta: float):
+    u(q, np.pi / 2, theta / 2 + np.pi / 2, np.pi / 2)
+
+
+@aq.gate
+def ry(q: aq.Qubit, theta: float):
+    u(q, np.pi, theta / 2 + np.pi, np.pi)
+
+
+@aq.gate
+def cnot(q0: aq.Qubit, q1: aq.Qubit):
+    ry(q0, np.pi / 2)
+    ms(q0, q1, 0, 0, np.pi / 2)
+    rx(q0, -np.pi / 2)
+    rx(q1, -np.pi / 2)
+    ry(q0, -np.pi / 2)

src/braket/experimental/autoqasm/api.py

@@ -28,6 +28,7 @@
 import braket.experimental.autoqasm.transpiler as aq_transpiler
 import braket.experimental.autoqasm.types as aq_types
 from braket.aws import AwsDevice
+from braket.devices.device import Device
 from braket.experimental.autoqasm import errors
 from braket.experimental.autoqasm.autograph.core import converter
 from braket.experimental.autoqasm.autograph.impl.api_core import (
@@ -41,7 +42,7 @@
 
 
 def main(
-    *args, num_qubits: Optional[int] = None, device: Optional[Union[AwsDevice, str]] = None
+    *args, num_qubits: Optional[int] = None, device: Optional[Union[Device, str]] = None
 ) -> Callable[[Any], aq_program.Program]:
     """Decorator that converts a function into a callable that returns
     a Program object containing the quantum program.
@@ -52,8 +53,8 @@ def main(
     Args:
         num_qubits (Optional[int]): Configuration to set the total number of qubits to declare in
             the program.
-        device (Optional[Union[AwsDevice, str]]): Configuration to set the target device for the
-            program. Can be either an AwsDevice object or a valid Amazon Braket device ARN.
+        device (Optional[Union[Device, str]]): Configuration to set the target device for the
+            program. Can be either an Device object or a valid Amazon Braket device ARN.
 
     Returns:
         Callable[[Any], Program]: A callable which returns the converted
@@ -63,8 +64,8 @@ def main(
         device = AwsDevice(device)
 
     return _function_wrapper(
-        args,
-        _convert_main,
+        *args,
+        converter_callback=_convert_main,
         converter_args={"user_config": aq_program.UserConfig(num_qubits=num_qubits, device=device)},
     )
 
@@ -77,7 +78,7 @@ def subroutine(*args) -> Callable[[Any], aq_program.Program]:
         Callable[[Any], Program]: A callable which returns the converted
         quantum program when called.
     """
-    return _function_wrapper(args, _convert_subroutine)
+    return _function_wrapper(*args, converter_callback=_convert_subroutine)
 
 
 def gate(*args) -> Callable[[Any], None]:
@@ -87,7 +88,7 @@ def gate(*args) -> Callable[[Any], None]:
         Callable[[Any],]: A callable which can be used as a custom gate inside an
         aq.function or inside another aq.gate.
     """
-    return _function_wrapper(args, _convert_gate)
+    return _function_wrapper(*args, converter_callback=_convert_gate)
 
 
 def gate_calibration(*args, implements: Callable, **kwargs) -> Callable[[], GateCalibration]:
@@ -104,20 +105,21 @@ def gate_calibration(*args, implements: Callable, **kwargs) -> Callable[[], Gate
         Callable[[], GateCalibration]: A callable to be added to a main program using
         `with_calibrations` method of the main program.
     """
-    converter_args = {"gate_function": implements, **kwargs}
-
-    return _function_wrapper(args, _convert_calibration, converter_args)
+    return _function_wrapper(
+        *args,
+        converter_callback=_convert_calibration,
+        converter_args={"gate_function": implements, **kwargs},
+    )
 
 
 def _function_wrapper(
-    args: Tuple[Any],
+    *args: Tuple[Any],
     converter_callback: Callable,
     converter_args: Optional[Dict[str, Any]] = None,
 ) -> Callable[[Any], aq_program.Program]:
     """Wrapping and conversion logic around the user function `f`.
 
     Args:
-        args (Tuple[Any]): The arguments to the decorated function.
         converter_callback (Callable): The function converter, e.g., _convert_main.
         converter_args (Optional[Dict[str, Any]]): Extra arguments for the function converter.
 
@@ -129,12 +131,11 @@ def _function_wrapper(
         # This the case where a decorator is called with only keyword args, for example:
         #     @aq.main(num_qubits=4)
         #     def my_function():
-        # To make this work, here we simply return another wrapper function which expects
-        # a Callable as the first argument.
-        def _function_wrapper_with_params(*args) -> Callable[[Any], aq_program.Program]:
-            return _function_wrapper(args, converter_callback, converter_args=converter_args)
-
-        return _function_wrapper_with_params
+        # To make this work, here we simply return a partial application of this function
+        # which still expects a Callable as the first argument.
+        return functools.partial(
+            _function_wrapper, converter_callback=converter_callback, converter_args=converter_args
+        )
 
     f = args[0]
     if is_autograph_artifact(f):

src/braket/experimental/autoqasm/errors.py

@@ -45,6 +45,22 @@ class InvalidCalibrationDefinition(AutoQasmError):
     """Calibration definition does not meet the necessary requirements."""
 
 
+class InvalidTargetQubit(AutoQasmError):
+    """Target qubit is invalid in the current context."""
+
+
+class UnsupportedGate(AutoQasmError):
+    """Gate is not supported by the target device."""
+
+
+class UnsupportedNativeGate(AutoQasmError):
+    """Native gate is not supported by the target device."""
+
+
+class VerbatimBlockNotAllowed(AutoQasmError):
+    """Verbatim block is not supported by the target device."""
+
+
 class UnknownQubitCountError(AutoQasmError):
     """Missing declaration for the number of qubits."""
 

src/braket/experimental/autoqasm/instructions/instructions.py

@@ -25,12 +25,11 @@
 def _qubit_instruction(
     name: str, qubits: List[QubitIdentifierType], *args: Any, is_unitary: bool = True
 ) -> None:
-    # If this is an instruction inside a gate definition, ensure that it only operates on
-    # qubits and angles which are passed as arguments to the gate definition.
     program_conversion_context = aq_program.get_program_conversion_context()
     program_conversion_context.validate_gate_targets(qubits, args)
 
     # Add the instruction to the program.
+    program_conversion_context.register_gate(name)
     program_mode = aq_program.ProgramMode.UNITARY if is_unitary else aq_program.ProgramMode.NONE
     oqpy_program = program_conversion_context.get_oqpy_program(mode=program_mode)
     oqpy_program.gate([_qubit(q) for q in qubits], name, *args)

src/braket/experimental/autoqasm/instructions/qubits.py

@@ -14,12 +14,14 @@
 
 """Utility functions that handle qubit construction and naming."""
 
+import re
 from functools import singledispatch
-from typing import Any, Union
+from typing import Any, List, Union
 
 import oqpy.base
 from openpulse.printer import dumps
 
+from braket.circuits.qubit_set import QubitSet
 from braket.experimental.autoqasm import constants, errors, program
 
 QubitIdentifierType = Union[int, oqpy._ClassicalVar, oqpy.base.OQPyExpression, str, oqpy.Qubit]
@@ -37,6 +39,26 @@ def is_qubit_identifier_type(qubit: Any) -> bool:
     return isinstance(qubit, QubitIdentifierType.__args__)
 
 
+def physical_qubit_to_braket_qubit(qids: List[str]) -> QubitSet:
+    """Convert a physical qubit label to a QubitSet.
+
+    Args:
+        qids (List[str]): Physical qubit labels.
+
+    Returns:
+        QubitSet: Represent physical qubits.
+    """
+    braket_qubits = []
+    for qid in qids:
+        if not (isinstance(qid, str) and re.match(r"\$\d+", qid)):
+            raise ValueError(
+                f"invalid physical qubit label: '{qid}'. Physical qubit must be labeled as a string"
+                "with `$` followed by an integer. For example: `$1`."
+            )
+        braket_qubits.append(int(qid[1:]))
+    return QubitSet(braket_qubits)
+
+
 def _global_qubit_register(qubit_idx_expr: Union[int, str]) -> str:
     # TODO: We should index into a oqpy.QubitArray rather
     # than manually generating the string to index into

src/braket/experimental/autoqasm/program/__init__.py

@@ -15,7 +15,7 @@
 for AutoQASM.
 """
 
-from .pragmas import Verbatim as verbatim  # noqa: F401
+from .pragmas import verbatim  # noqa: F401
 from .program import (  # noqa: F401
     GateArgs,
     Program,

src/braket/experimental/autoqasm/program/pragmas.py

@@ -20,7 +20,7 @@
 
     @aq.main
     def pragma_example() -> None:
-        with aq.Verbatim():
+        with aq.verbatim():
             h(0)
             cnot(0, 1)
         x(0)
@@ -29,24 +29,21 @@ def pragma_example() -> None:
 """
 
 
-import oqpy.base
+import contextlib
 
 from braket.experimental.autoqasm import program
 
 
-class Verbatim:
+@contextlib.contextmanager
+def verbatim() -> None:
     """Context management protocol that, when used with a `with` statement, wraps the code block
-    in a verbatim box.
+    in a verbatim block.
 
-    The verbatim pragma around a code block specifies that operations are to be executed as
+    A verbatim block specifies that operations contained within the block are to be executed as
     programmed without compilation or modification of any sort.
-    """
-
-    def __enter__(self):
-        oqpy_program = program.get_program_conversion_context().get_oqpy_program()
-        self.box = oqpy.Box(oqpy_program)
-        oqpy_program.pragma("braket verbatim")
-        self.box.__enter__()
 
-    def __exit__(self, exc_type, exc, traceback):
-        return self.box.__exit__(exc_type, exc, traceback)
+    Raises:
+        errors.VerbatimBlockNotAllowed: If the target device does not support verbatim blocks.
+    """
+    with program.get_program_conversion_context().verbatim_block():
+        yield