[Frontend] Make tests toml-schema independent

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/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,

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)

frontend/catalyst/third_party/cuda/catalyst_to_cuda_interpreter.py

@@ -39,6 +39,7 @@
 import pennylane as qml
 from jax.tree_util import tree_unflatten
 
+from catalyst.device import BackendInfo
 from catalyst.jax_primitives import (
     AbstractObs,
     adjoint_p,
@@ -75,7 +76,6 @@
 from catalyst.qfunc import QFunc
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.patching import Patcher
-from catalyst.utils.runtime import BackendInfo
 
 from .primitives import (
     cuda_inst,
@@ -820,7 +820,7 @@ def get_jaxpr(self, *args):
             an MLIR module
         """
 
-        def cudaq_backend_info(device, _config) -> BackendInfo:
+        def cudaq_backend_info(device, _capabilities) -> BackendInfo:
             """The extract_backend_info should not be run by the cuda compiler as it is
             catalyst-specific. We need to make this API a bit nicer for third-party compilers.
             """

frontend/catalyst/third_party/cuda/primitives/__init__.py

@@ -24,7 +24,7 @@
 
 # We disable protected access in particular to avoid warnings with cudaq._pycuda.
 # And we disable unused-argument to avoid unused arguments in abstract_eval, particularly kwargs.
-# pylint: disable=protected-access,unused-argument,abstract-method,line-too-long
+# pylint: disable=protected-access,unused-argument,line-too-long
 
 
 class AbsCudaQState(jax.core.AbstractValue):

frontend/catalyst/utils/paths.py

@@ -0,0 +1,38 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Utility code for keeping paths
+"""
+import os
+import os.path
+
+from catalyst._configuration import INSTALLED
+
+package_root = os.path.dirname(__file__)
+
+# Default paths to dep libraries
+DEFAULT_LIB_PATHS = {
+    "llvm": os.path.join(package_root, "../../../mlir/llvm-project/build/lib"),
+    "runtime": os.path.join(package_root, "../../../runtime/build/lib"),
+    "enzyme": os.path.join(package_root, "../../../mlir/Enzyme/build/Enzyme"),
+    "oqc_runtime": os.path.join(package_root, "../../catalyst/third_party/oqc/src/build"),
+}
+
+
+def get_lib_path(project, env_var):
+    """Get the library path."""
+    if INSTALLED:
+        return os.path.join(package_root, "..", "lib")  # pragma: no cover
+    return os.getenv(env_var, DEFAULT_LIB_PATHS.get(project, ""))

frontend/catalyst/utils/runtime.py

@@ -124,6 +124,7 @@ def validate_config_with_device(device: qml.QubitDevice, config: TOMLDocument) -
 
     Raises: CompileError
     """
+    # pylint: disable=too-many-function-args
 
     if not config["compilation"]["qjit_compatible"]:
         raise CompileError(