Native quantum control toml (#554)

By this PR we change the TOML file schema in a way that allows us to
represent quantum control at per-gate level.

[sc-57159]
[sc-57172]

---------

Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: Romain Moyard <rmoyard@gmail.com>

.gitignore

@@ -13,6 +13,7 @@ pennylane_catalyst.egg-info
 dist
 frontend/catalyst/bin
 frontend/catalyst/lib
+frontend/catalyst/_revision.py
 frontend/mlir_quantum
 
 # IDEs

Makefile

@@ -17,6 +17,7 @@ ENABLE_OPENQASM?=ON
 TEST_BACKEND ?= "lightning.qubit"
 TEST_BRAKET ?= NONE
 ENABLE_ASAN ?= OFF
+TOML_SPECS ?= $(shell find ./runtime ./frontend -name '*.toml')
 
 PLATFORM := $(shell uname -s)
 ifeq ($(PLATFORM),Linux)
@@ -108,7 +109,10 @@ dummy_device:
 	$(MAKE) -C runtime dummy_device
 
 .PHONY: test test-runtime test-frontend lit pytest test-demos
-test: test-runtime test-frontend test-demos
+test: test-runtime test-frontend test-demos test-toml-spec
+
+test-toml-spec:
+	$(PYTHON) ./bin/toml-check.py $(TOML_SPECS)
 
 test-runtime:
 	$(MAKE) -C runtime test

bin/toml-check.py

@@ -0,0 +1,94 @@
+#!/usr/bin/python3
+# 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.
+
+"""
+This program checks the syntax of quantum device configuration files. It is a strict parser of
+TOML format, narrowed down to match our requirements.  For the Lark's EBNF dialect syntax, see the
+Lark grammar reference:
+ * https://lark-parser.readthedocs.io/en/latest/grammar.html
+"""
+
+import sys
+from argparse import ArgumentParser
+from textwrap import dedent
+
+try:
+    from lark import Lark, LarkError, UnexpectedInput
+except ImportError as e:
+    raise RuntimeError(
+        "toml-check.py requires `lark` library. Consider using `pip install lark`"
+    ) from e
+
+parser = Lark(
+    dedent(
+        """
+        start: schema_body \
+               gates_native_section \
+               gates_decomp_section \
+               gates_matrix_section \
+               gates_observables_section \
+               measurement_processes_section \
+               compilation_section \
+               options_section?
+        schema_body: schema_decl
+        gates_native_section: "[operators.gates.native]" gate_decls
+        gates_decomp_section: "[operators.gates.decomp]" gate_decls
+        gates_matrix_section: "[operators.gates.matrix]" gate_decls
+        gates_observables_section: "[operators.observables]" gate_decls
+        measurement_processes_section: "[measurement_processes]" gate_decls
+        compilation_section: "[compilation]" flag_decl*
+        options_section: "[options]" option_decl*
+        schema_decl: "schema" "=" "2"
+        gate_decls: (gate_decl)*
+        gate_decl: name "=" "{" (gate_trait ("," gate_trait)*)? "}"
+        gate_trait: gate_condition | gate_properties
+        gate_condition: "condition" "=" "[" ( "\\"finiteshots\\"" | "\\"analytic\\"" ) "]"
+        gate_properties: "properties" "=" "[" gate_property ("," gate_property)* "]"
+        gate_property: "\\"controllable\\"" | "\\"invertible\\"" | "\\"differentiable\\""
+        flag_decl: ( "qjit_compatible" | "runtime_code_generation" | \
+                     "mid_circuit_measurement" | "dynamic_qubit_management" ) "=" boolean
+        option_decl: name "=" (name | "\\"" name "\\"")
+        name: /[a-zA-Z0-9_]+/
+        boolean: "true" | "false"
+        COMMENT: "#" /./*
+        %import common.WS
+        %ignore WS
+        %ignore COMMENT
+        """
+    )
+)
+
+
+if __name__ == "__main__":
+    ap = ArgumentParser(prog="toml-check.py")
+    ap.add_argument(
+        "filenames", metavar="TOML", type=str, nargs="+", help="One or more *toml files to check"
+    )
+    ap.add_argument("--verbose", action="store_true", help="Be verbose")
+    fname = None
+    try:
+        arguments = ap.parse_args(sys.argv[1:])
+        for fname in arguments.filenames:
+            with open(fname, "r", encoding="utf-8") as f:
+                contents = f.read()
+            tree = parser.parse(contents)
+            if arguments.verbose:
+                print(tree.pretty())
+    except UnexpectedInput as e:
+        print(f"toml-check: error in {fname}:{e.line}:{e.column}", file=sys.stderr)
+        raise e
+    except LarkError as e:
+        print(f"toml-check: error in {fname}", file=sys.stderr)
+        raise e

doc/changelog.md

@@ -9,6 +9,11 @@
 
 <h3>Improvements</h3>
 
+* An updated quantum device specification format is now supported by Catalyst. The toml schema 2
+  configs allow device autors to specify individual gate properties such as native quantum control
+  support, gate invertibility or differentiability.
+  [(#554)](https://github.com/PennyLaneAI/catalyst/pull/554)
+
 * Catalyst now supports devices built from the 
   [new PennyLane device API](https://docs.pennylane.ai/en/stable/code/api/pennylane.devices.Device.html).
   [(#565)](https://github.com/PennyLaneAI/catalyst/pull/565)
@@ -57,8 +62,9 @@ This release contains contributions from (in alphabetical order):
 
 Ali Asadi,
 David Ittah,
+Erick Ochoa Lopez,
 Romain Moyard,
-Erick Ochoa Lopez.
+Sergei Mironov.
 
 # Release 0.5.0
 

doc/dev/custom_devices.rst

@@ -230,130 +230,124 @@ headers and fields are generally required, unless stated otherwise.
 .. code-block:: toml
 
         # Which version of the specification format is being used.
-        schema = 1
-
-        [device]
-        name = "dummy.device.qubit"
-
-        [operators]
-        # Observables supported by the device
-        observables = [
-                "PauliX",
-                "PauliY",
-                "PauliZ",
-                "Hadamard",
-                "Hermitian",
-                "Identity",
-                "Projector",
-                "SparseHamiltonian",
-                "Hamiltonian",
-                "Sum",
-                "SProd",
-                "Prod",
-                "Exp",
-        ]
+        schema = 2
 
         # The union of all gate types listed in this section must match what
         # the device considers "supported" through PennyLane's device API.
-        [[operators.gates]]
-        native = [
-                # Operators that shouldn't be decomposed.
-                "QubitUnitary",
-                "PauliX",
-                "PauliY",
-                "PauliZ",
-                "MultiRZ",
-                "Hadamard",
-                "S",
-                "T",
-                "CNOT",
-                "SWAP",
-                "CSWAP",
-                "Toffoli",
-                "CY",
-                "CZ",
-                "PhaseShift",
-                "ControlledPhaseShift",
-                "RX",
-                "RY",
-                "RZ",
-                "Rot",
-                "CRX",
-                "CRY",
-                "CRZ",
-                "CRot",
-                "Identity",
-                "IsingXX",
-                "IsingYY",
-                "IsingZZ",
-                "IsingXY",
-        ]
+        # The gate definition has the following format:
+        #
+        #   GATE = { properties = [ PROPS ], condition = [ COND ] }
+        #
+        # Where:
+        #
+        #   PROPS: zero or more comma-separated quoted strings:
+        #          "controllable", "invertible", "differentiable"
+        #   COND: quoted string, on of:
+        #         "analytic", "finiteshots"
+        #
+        [operators.gates.native]
+
+        QubitUnitary = { properties = [ "controllable", "invertible"]  }
+        PauliX = { properties = [ "controllable", "invertible"] }
+        PauliY = { properties = [ "controllable", "invertible"] }
+        PauliZ = { properties = [ "controllable", "invertible"] }
+        MultiRZ = { properties = [ "controllable", "invertible" ] }
+        Hadamard = { properties = [ "controllable", "invertible"] }
+        S = { properties = [ "controllable", "invertible" ] }
+        T = { properties = [ "controllable", "invertible" ] }
+        CNOT = { properties = [ "invertible" ] }
+        SWAP = { properties = [ "controllable", "invertible" ] }
+        CSWAP = { properties = [ "invertible" ] }
+        Toffoli = { properties = [ "controllable", "invertible" ] }
+        CY = { properties = [ "invertible" ] }
+        CZ = { properties = [ "invertible" ] }
+        PhaseShift = { properties = [ "controllable", "invertible" ] }
+        ControlledPhaseShift = { properties = [ "controllable", "invertible" ] }
+        RX = { properties = [ "controllable", "invertible" ] }
+        RY = { properties = [ "controllable", "invertible" ] }
+        RZ = { properties = [ "controllable", "invertible" ] }
+        Rot = { properties = [ "controllable", "invertible" ] }
+        CRX = { properties = [ "invertible" ] }
+        CRY = { properties = [ "invertible" ] }
+        CRZ = { properties = [ "invertible" ] }
+        CRot = { properties = [ "invertible" ] }
+        Identity = { properties = [ "controllable", "invertible" ] }
+        IsingXX = { properties = [ "controllable", "invertible" ] }
+        IsingYY = { properties = [ "controllable", "invertible" ] }
+        IsingZZ = { properties = [ "controllable", "invertible" ] }
+        IsingXY = { properties = [ "controllable", "invertible" ] }
 
         # Operators that should be decomposed according to the algorithm used
         # by PennyLane's device API.
         # Optional, since gates not listed in this list will typically be decomposed by
         # default, but can be useful to express a deviation from this device's regular
         # strategy in PennyLane.
-        decomp = [
-                "SX",
-                "ISWAP",
-                "PSWAP",
-                "SISWAP",
-                "SQISW",
-                "CPhase",
-                "BasisState",
-                "QubitStateVector",
-                "StatePrep",
-                "ControlledQubitUnitary",
-                "DiagonalQubitUnitary",
-                "SingleExcitation",
-                "SingleExcitationPlus",
-                "SingleExcitationMinus",
-                "DoubleExcitation",
-                "DoubleExcitationPlus",
-                "DoubleExcitationMinus",
-                "QubitCarry",
-                "QubitSum",
-                "OrbitalRotation",
-                "QFT",
-                "ECR",
-        ]
+        [operators.gates.decomp]
+
+        SX = {}
+        ISWAP = {}
+        PSWAP = {}
+        SISWAP = {}
+        SQISW = {}
+        CPhase = {}
+        BasisState = {}
+        QubitStateVector = {}
+        StatePrep = {}
+        ControlledQubitUnitary = {}
+        DiagonalQubitUnitary = {}
+        SingleExcitation = {}
+        SingleExcitationPlus = {}
+        SingleExcitationMinus = {}
+        DoubleExcitation = {}
+        DoubleExcitationPlus = {}
+        DoubleExcitationMinus = {}
+        QubitCarry = {}
+        QubitSum = {}
+        OrbitalRotation = {}
+        QFT = {}
+        ECR = {}
 
         # Gates which should be translated to QubitUnitary
-        matrix = [
-                "MultiControlledX",
-        ]
+        [operators.gates.matrix]
+
+        MultiControlledX = {}
+
+        # Observables supported by the device
+        [operators.observables]
+
+        PauliX = {}
+        PauliY = {}
+        PauliZ = {}
+        Hadamard = {}
+        Hermitian = {}
+        Identity = {}
+        Projector = {}
+        SparseHamiltonian = {}
+        Hamiltonian = {}
+        Sum = {}
+        SProd = {}
+        Prod = {}
+        Exp = {}
 
         [measurement_processes]
-        exactshots = [
-                "Expval",
-                "Var",
-                "Probs",
-                "State",
-        ]
-        finiteshots = [
-                "Expval",
-                "Var",
-                "Probs",
-                "Sample",
-                "Counts",
-        ]
+
+        Expval = {}
+        Var = {}
+        Probs = {}
+        Sample = {}
+        Count = { condition = [ "finiteshots" ] }
 
         [compilation]
+
         # If the device is compatible with qjit
         qjit_compatible = true
         # If the device requires run time generation of the quantum circuit.
         runtime_code_generation = false
-        # If the device supports adjoint
-        quantum_adjoint = true
-        # If the device supports quantum control instructions natively
-        quantum_control = false
         # If the device supports mid circuit measurements natively
         mid_circuit_measurement = true
-
         # This field is currently unchecked but it is reserved for the purpose of
         # determining if the device supports dynamic qubit allocation/deallocation.
-        dynamic_qubit_management = false 
+        dynamic_qubit_management = false
 
         [options]
         # Options is an optional field.

frontend/catalyst/cuda/catalyst_to_cuda_interpreter.py

@@ -76,7 +76,7 @@
 from catalyst.pennylane_extensions import QFunc
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.patching import Patcher
-from catalyst.utils.toml import toml_load
+from catalyst.utils.runtime import BackendInfo
 
 from .primitives import (
     cuda_inst,
@@ -821,14 +821,12 @@ def get_jaxpr(self, *args):
             an MLIR module
         """
 
-        def cudaq_backend_info(device):
+        def cudaq_backend_info(device, _config) -> 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.
             """
-            with open(device.config, "rb") as f:
-                config = toml_load(f)
-
-            return config, device.name, None, None
+            device_name = device.short_name if isinstance(device, qml.Device) else device.name
+            return BackendInfo(device_name, device.name, "", {})
 
         with Patcher(
             (catalyst.pennylane_extensions.QFunc, "extract_backend_info", cudaq_backend_info),