[Frontend] Pass capabilities to the decomposer

frontend/catalyst/compiler.py

@@ -33,7 +33,7 @@
 
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.filesystem import Directory
-from catalyst.utils.runtime import get_lib_path
+from catalyst.utils.paths import get_lib_path
 
 package_root = os.path.dirname(__file__)
 

frontend/catalyst/device/__init__.py

@@ -16,9 +16,11 @@
 Internal API for the device module.
 """
 
-from catalyst.device.qjit_device import QJITDevice, QJITDeviceNewAPI
-
-__all__ = (
-    "QJITDevice",
-    "QJITDeviceNewAPI",
+from catalyst.device.qjit_device import (
+    BackendInfo,
+    QJITDevice,
+    QJITDeviceNewAPI,
+    extract_backend_info,
 )
+
+__all__ = ("QJITDevice", "QJITDeviceNewAPI", "BackendInfo", "extract_backend_info")

frontend/catalyst/device/decomposition.py

@@ -17,6 +17,8 @@
 compilation & execution on devices.
 """
 
+from functools import partial
+
 import jax
 import pennylane as qml
 from pennylane import transform
@@ -37,15 +39,16 @@
 from catalyst.jax_tracer import HybridOpRegion, has_nested_tapes
 from catalyst.tracing.contexts import EvaluationContext
 from catalyst.utils.exceptions import CompileError
+from catalyst.utils.toml import DeviceCapabilities
 
 
-def catalyst_decomposer(op):
+def catalyst_decomposer(op, capabilities: DeviceCapabilities):
     """A decomposer for catalyst, to be passed to the decompose transform. Takes an operator and
     returns the default decomposition, unless the operator should decompose to a QubitUnitary.
     Raises a CompileError for MidMeasureMP"""
     if isinstance(op, MidMeasureMP):
         raise CompileError("Must use 'measure' from Catalyst instead of PennyLane.")
-    if op.name in {"MultiControlledX", "BlockEncode"} or isinstance(op, qml.ops.Controlled):
+    if capabilities.to_matrix_ops.get(op.name) or isinstance(op, qml.ops.Controlled):
         return _decompose_to_matrix(op)
     return op.decomposition()
 
@@ -55,7 +58,7 @@ def catalyst_decompose(
     tape: qml.tape.QuantumTape,
     ctx,
     stopping_condition,
-    decomposer=catalyst_decomposer,
+    capabilities,
     max_expansion=None,
 ):
     """Decompose operations until the stopping condition is met.
@@ -76,7 +79,7 @@ def catalyst_decompose(
         tape,
         stopping_condition,
         skip_initial_state_prep=False,
-        decomposer=decomposer,
+        decomposer=partial(catalyst_decomposer, capabilities=capabilities),
         max_expansion=max_expansion,
         name="catalyst on this device",
         error=CompileError,
@@ -85,7 +88,7 @@ def catalyst_decompose(
     new_ops = []
     for op in toplevel_tape.operations:
         if has_nested_tapes(op):
-            op = _decompose_nested_tapes(op, ctx, stopping_condition, decomposer, max_expansion)
+            op = _decompose_nested_tapes(op, ctx, stopping_condition, capabilities, max_expansion)
         new_ops.append(op)
     tape = qml.tape.QuantumScript(new_ops, tape.measurements, shots=tape.shots)
 
@@ -103,7 +106,7 @@ def _decompose_to_matrix(op):
     return [op]
 
 
-def _decompose_nested_tapes(op, ctx, stopping_condition, decomposer, max_expansion):
+def _decompose_nested_tapes(op, ctx, stopping_condition, capabilities, max_expansion):
     new_regions = []
     for region in op.regions:
         if region.quantum_tape is None:
@@ -114,7 +117,7 @@ def _decompose_nested_tapes(op, ctx, stopping_condition, decomposer, max_expansi
                     region.quantum_tape,
                     ctx=ctx,
                     stopping_condition=stopping_condition,
-                    decomposer=decomposer,
+                    capabilities=capabilities,
                     max_expansion=max_expansion,
                 )
                 new_tape = tapes[0]

frontend/catalyst/device/qjit_device.py

@@ -16,10 +16,13 @@
 This module contains device stubs for the old and new PennyLane device API, which facilitate
 the application of decomposition and other device pre-processing routines.
 """
-
+import os
+import pathlib
+import platform
 from copy import deepcopy
+from dataclasses import dataclass
 from functools import partial
-from typing import Optional, Set
+from typing import Any, Dict, Optional, Set
 
 import pennylane as qml
 from pennylane.measurements import MidMeasureMP
@@ -32,13 +35,10 @@
 )
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.patching import Patcher
-from catalyst.utils.runtime import BackendInfo, device_get_toml_config
+from catalyst.utils.paths import get_lib_path
 from catalyst.utils.toml import (
     DeviceCapabilities,
     OperationProperties,
-    ProgramFeatures,
-    TOMLDocument,
-    get_device_capabilities,
     intersect_operations,
     pennylane_operation_set,
 )
@@ -84,6 +84,87 @@
     for op in RUNTIME_OPERATIONS
 }
 
+# TODO: This should be removed after implementing `get_c_interface`
+# for the following backend devices:
+SUPPORTED_RT_DEVICES = {
+    "lightning.qubit": ("LightningSimulator", "librtd_lightning"),
+    "lightning.kokkos": ("LightningKokkosSimulator", "librtd_lightning"),
+    "braket.aws.qubit": ("OpenQasmDevice", "librtd_openqasm"),
+    "braket.local.qubit": ("OpenQasmDevice", "librtd_openqasm"),
+}
+
+
+@dataclass
+class BackendInfo:
+    """Backend information"""
+
+    device_name: str
+    c_interface_name: str
+    lpath: str
+    kwargs: Dict[str, Any]
+
+
+def extract_backend_info(device: qml.QubitDevice, capabilities: DeviceCapabilities) -> BackendInfo:
+    """Extract the backend info from a quantum device. The device is expected to carry a reference
+    to a valid TOML config file."""
+    # pylint: disable=too-many-branches
+
+    dname = device.name
+    if isinstance(device, qml.Device):
+        dname = device.short_name
+
+    device_name = ""
+    device_lpath = ""
+    device_kwargs = {}
+
+    if dname in SUPPORTED_RT_DEVICES:
+        # Support backend devices without `get_c_interface`
+        device_name = SUPPORTED_RT_DEVICES[dname][0]
+        device_lpath = get_lib_path("runtime", "RUNTIME_LIB_DIR")
+        sys_platform = platform.system()
+
+        if sys_platform == "Linux":
+            device_lpath = os.path.join(device_lpath, SUPPORTED_RT_DEVICES[dname][1] + ".so")
+        elif sys_platform == "Darwin":  # pragma: no cover
+            device_lpath = os.path.join(device_lpath, SUPPORTED_RT_DEVICES[dname][1] + ".dylib")
+        else:  # pragma: no cover
+            raise NotImplementedError(f"Platform not supported: {sys_platform}")
+    elif hasattr(device, "get_c_interface"):
+        # Support third party devices with `get_c_interface`
+        device_name, device_lpath = device.get_c_interface()
+    else:
+        raise CompileError(f"The {dname} device does not provide C interface for compilation.")
+
+    if not pathlib.Path(device_lpath).is_file():
+        raise CompileError(f"Device at {device_lpath} cannot be found!")
+
+    if hasattr(device, "shots"):
+        if isinstance(device, qml.Device):
+            device_kwargs["shots"] = device.shots if device.shots else 0
+        else:
+            # TODO: support shot vectors
+            device_kwargs["shots"] = device.shots.total_shots if device.shots else 0
+
+    if dname == "braket.local.qubit":  # pragma: no cover
+        device_kwargs["device_type"] = dname
+        device_kwargs["backend"] = (
+            # pylint: disable=protected-access
+            device._device._delegate.DEVICE_ID
+        )
+    elif dname == "braket.aws.qubit":  # pragma: no cover
+        device_kwargs["device_type"] = dname
+        device_kwargs["device_arn"] = device._device._arn  # pylint: disable=protected-access
+        if device._s3_folder:  # pylint: disable=protected-access
+            device_kwargs["s3_destination_folder"] = str(
+                device._s3_folder  # pylint: disable=protected-access
+            )
+
+    for k, v in capabilities.options.items():
+        if hasattr(device, v):
+            device_kwargs[k] = getattr(device, v)
+
+    return BackendInfo(dname, device_name, device_lpath, device_kwargs)
+
 
 def get_qjit_device_capabilities(target_capabilities: DeviceCapabilities) -> Set[str]:
     """Calculate the set of supported quantum gates for the QJIT device from the gates
@@ -165,7 +246,7 @@
 
     def __init__(
         self,
-        target_config: TOMLDocument,
+        original_device_capabilities: DeviceCapabilities,
         shots=None,
         wires=None,
         backend: Optional[BackendInfo] = None,
@@ -175,23 +256,18 @@
         self.backend_name = backend.c_interface_name if backend else "default"
         self.backend_lib = backend.lpath if backend else ""
         self.backend_kwargs = backend.kwargs if backend else {}
-        device_name = backend.device_name if backend else "default"
 
-        program_features = ProgramFeatures(shots is not None)
-        target_device_capabilities = get_device_capabilities(
-            target_config, program_features, device_name
-        )
-        self.capabilities = get_qjit_device_capabilities(target_device_capabilities)
+        self.qjit_capabilities = get_qjit_device_capabilities(original_device_capabilities)
 
     @property
     def operations(self) -> Set[str]:
         """Get the device operations using PennyLane's syntax"""
-        return pennylane_operation_set(self.capabilities.native_ops)
+        return pennylane_operation_set(self.qjit_capabilities.native_ops)
 
     @property
     def observables(self) -> Set[str]:
         """Get the device observables"""
-        return pennylane_operation_set(self.capabilities.native_obs)
+        return pennylane_operation_set(self.qjit_capabilities.native_obs)
 
     def apply(self, operations, **kwargs):
         """
@@ -270,6 +346,7 @@
     def __init__(
         self,
         original_device,
+        original_device_capabilities: DeviceCapabilities,
         backend: Optional[BackendInfo] = None,
     ):
         self.original_device = original_device
@@ -285,29 +362,23 @@
         self.backend_name = backend.c_interface_name if backend else "default"
         self.backend_lib = backend.lpath if backend else ""
         self.backend_kwargs = backend.kwargs if backend else {}
-        device_name = backend.device_name if backend else "default"
 
-        target_config = device_get_toml_config(original_device)
-        program_features = ProgramFeatures(original_device.shots is not None)
-        target_device_capabilities = get_device_capabilities(
-            target_config, program_features, device_name
-        )
-        self.capabilities = get_qjit_device_capabilities(target_device_capabilities)
+        self.qjit_capabilities = get_qjit_device_capabilities(original_device_capabilities)
 
     @property
     def operations(self) -> Set[str]:
         """Get the device operations"""
-        return pennylane_operation_set(self.capabilities.native_ops)
+        return pennylane_operation_set(self.qjit_capabilities.native_ops)
 
     @property
     def observables(self) -> Set[str]:
         """Get the device observables"""
-        return pennylane_operation_set(self.capabilities.native_obs)
+        return pennylane_operation_set(self.qjit_capabilities.native_obs)
 
     @property
     def measurement_processes(self) -> Set[str]:
         """Get the device measurement processes"""
-        return self.capabilities.measurement_processes
+        return self.qjit_capabilities.measurement_processes
 
     def preprocess(
         self,
@@ -321,7 +392,12 @@
         program = TransformProgram()
 
         ops_acceptance = partial(catalyst_acceptance, operations=self.operations)
-        program.add_transform(catalyst_decompose, ctx=ctx, stopping_condition=ops_acceptance)
+        program.add_transform(
+            catalyst_decompose,
+            ctx=ctx,
+            stopping_condition=ops_acceptance,
+            capabilities=self.qjit_capabilities,
+        )
 
         if self.measurement_processes == {"Counts"}:
             program.add_transform(measurements_from_counts)

frontend/catalyst/jax_primitives.py

@@ -76,7 +76,7 @@
 from catalyst.utils.extra_bindings import FromElementsOp, TensorExtractOp
 from catalyst.utils.types import convert_shaped_arrays_to_tensors
 
-# pylint: disable=unused-argument,abstract-method,too-many-lines
+# pylint: disable=unused-argument,too-many-lines
 
 #########
 # Types #

frontend/catalyst/qfunc.py

@@ -22,21 +22,25 @@
 from jax.core import eval_jaxpr
 from jax.tree_util import tree_flatten, tree_unflatten
 
-from catalyst.device import QJITDevice, QJITDeviceNewAPI
+from catalyst.device import (
+    BackendInfo,
+    QJITDevice,
+    QJITDeviceNewAPI,
+    extract_backend_info,
+)
 from catalyst.jax_extras import (
     deduce_avals,
     get_implicit_and_explicit_flat_args,
     unzip2,
 )
 from catalyst.jax_primitives import func_p
 from catalyst.jax_tracer import trace_quantum_function
-from catalyst.utils.runtime import (
-    BackendInfo,
-    device_get_toml_config,
-    extract_backend_info,
-    validate_config_with_device,
+from catalyst.utils.toml import (
+    DeviceCapabilities,
+    ProgramFeatures,
+    get_device_capabilities,
+    validate_device_capabilities,
 )
-from catalyst.utils.toml import TOMLDocument
 
 
 class QFunc:
@@ -54,26 +58,34 @@ def __new__(cls):
         raise NotImplementedError()  # pragma: no-cover
 
     @staticmethod
-    def extract_backend_info(device: qml.QubitDevice, config: TOMLDocument) -> BackendInfo:
+    def extract_backend_info(
+        device: qml.QubitDevice, capabilities: DeviceCapabilities
+    ) -> BackendInfo:
         """Wrapper around extract_backend_info in the runtime module."""
-        return extract_backend_info(device, config)
+        return extract_backend_info(device, capabilities)
 
-    # pylint: disable=no-member
+    # pylint: disable=no-member, attribute-defined-outside-init
     def __call__(self, *args, **kwargs):
         assert isinstance(self, qml.QNode)
 
-        config = device_get_toml_config(self.device)
-        validate_config_with_device(self.device, config)
-        backend_info = QFunc.extract_backend_info(self.device, config)
+        device = self.device
+        program_features = ProgramFeatures(device.shots is not None)
+        device_capabilities = get_device_capabilities(device, program_features)
+        backend_info = QFunc.extract_backend_info(device, device_capabilities)
+
+        # Validate decive operations against the declared capabilities
+        validate_device_capabilities(device, device_capabilities)
 
         if isinstance(self.device, qml.devices.Device):
-            device = QJITDeviceNewAPI(self.device, backend_info)
+            self.qjit_device = QJITDeviceNewAPI(device, device_capabilities, backend_info)
         else:
-            device = QJITDevice(config, self.device.shots, self.device.wires, backend_info)
+            self.qjit_device = QJITDevice(
+                device_capabilities, device.shots, device.wires, backend_info
+            )
 
         def _eval_quantum(*args):
             closed_jaxpr, out_type, out_tree = trace_quantum_function(
-                self.func, device, args, kwargs, qnode=self
+                self.func, self.qjit_device, args, kwargs, qnode=self
             )
             args_expanded = get_implicit_and_explicit_flat_args(None, *args)
             res_expanded = eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.consts, *args_expanded)