[Frontend] Pass capabilities to the decomposer (#749)

**Context:** Catalyst needs to deal with operations which are not
supported by the target device. This PR addresses
decomposition-to-matrix branch.

**Description of the Change:** We read gates to be decomposed to matrix
from the device capabilities (which in turn are described in the device
toml file).
[sc-62011]

**Benefits:** Device authors gain control over the decomposition
strategies

**Possible Drawbacks:**

**Related GitHub Issues:**
Requires #712

---------

Co-authored-by: lillian542 <[email protected]>

Author: Sergei Mironov

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

@@ -393,7 +393,12 @@ def preprocess(
         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/test/lit/test_decomposition.py

@@ -15,27 +15,29 @@
 # RUN: %PYTHON %s | FileCheck %s
 # pylint: disable=line-too-long
 
+import platform
 from copy import deepcopy
 
 import jax
 import pennylane as qml
 
 from catalyst import cond, for_loop, measure, qjit, while_loop
+from catalyst.compiler import get_lib_path
 from catalyst.utils.toml import (
+    OperationProperties,
     ProgramFeatures,
     get_device_capabilities,
     pennylane_operation_set,
 )
 
 
-def get_custom_device_without(num_wires, discards):
+def get_custom_device_without(num_wires, discards=frozenset(), force_matrix=frozenset()):
     """Generate a custom device without gates in discards."""
 
-    class CustomDevice(qml.QubitDevice):
+    class CustomDevice(qml.devices.Device):
         """Custom Gate Set Device"""
 
         name = "Custom Device"
-        short_name = "lightning.qubit"
         pennylane_requires = "0.35.0"
         version = "0.0.2"
         author = "Tester"
@@ -55,12 +57,13 @@ def __init__(self, shots=None, wires=None):
             )
             custom_capabilities = deepcopy(lightning_capabilities)
             for gate in discards:
-                if gate in custom_capabilities.native_ops:
-                    custom_capabilities.native_ops.pop(gate)
-                if gate in custom_capabilities.to_decomp_ops:
-                    custom_capabilities.to_decomp_ops.pop(gate)
-                if gate in custom_capabilities.to_matrix_ops:
-                    custom_capabilities.to_matrix_ops.pop(gate)
+                custom_capabilities.native_ops.pop(gate, None)
+                custom_capabilities.to_decomp_ops.pop(gate, None)
+                custom_capabilities.to_matrix_ops.pop(gate, None)
+            for gate in force_matrix:
+                custom_capabilities.native_ops.pop(gate, None)
+                custom_capabilities.to_decomp_ops.pop(gate, None)
+                custom_capabilities.to_matrix_ops[gate] = OperationProperties(False, False, False)
             self.qjit_capabilities = custom_capabilities
 
         def apply(self, operations, **kwargs):
@@ -81,12 +84,27 @@ def observables(self):
             """Return PennyLane observables"""
             return pennylane_operation_set(self.qjit_capabilities.native_obs)
 
+        @staticmethod
+        def get_c_interface():
+            """Returns a tuple consisting of the device name, and
+            the location to the shared object with the C/C++ device implementation.
+            """
+            system_extension = ".dylib" if platform.system() == "Darwin" else ".so"
+            lib_path = (
+                get_lib_path("runtime", "RUNTIME_LIB_DIR") + "/librtd_dummy" + system_extension
+            )
+            return "dummy.remote", lib_path
+
+        def execute(self, circuits, execution_config):
+            """Execution."""
+            return circuits, execution_config
+
     return CustomDevice(wires=num_wires)
 
 
 def test_decompose_multicontrolledx():
     """Test decomposition of MultiControlledX."""
-    dev = get_custom_device_without(5, {"MultiControlledX"})
+    dev = get_custom_device_without(5, discards={"MultiControlledX"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -107,7 +125,7 @@ def decompose_multicontrolled_x1(theta: float):
 
 def test_decompose_multicontrolledx_in_conditional():
     """Test decomposition of MultiControlledX in conditional."""
-    dev = get_custom_device_without(5, {"MultiControlledX"})
+    dev = get_custom_device_without(5, discards={"MultiControlledX"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -133,7 +151,7 @@ def cond_fn():
 
 def test_decompose_multicontrolledx_in_while_loop():
     """Test decomposition of MultiControlledX in while loop."""
-    dev = get_custom_device_without(5, {"MultiControlledX"})
+    dev = get_custom_device_without(5, discards={"MultiControlledX"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -160,7 +178,7 @@ def loop(v):
 
 def test_decompose_multicontrolledx_in_for_loop():
     """Test decomposition of MultiControlledX in for loop."""
-    dev = get_custom_device_without(5, {"MultiControlledX"})
+    dev = get_custom_device_without(5, discards={"MultiControlledX"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -186,7 +204,7 @@ def loop(_):
 
 def test_decompose_rot():
     """Test decomposition of Rot gate."""
-    dev = get_custom_device_without(1, {"Rot", "C(Rot)"})
+    dev = get_custom_device_without(1, discards={"Rot", "C(Rot)"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -216,7 +234,7 @@ def decompose_rot(phi: float, theta: float, omega: float):
 
 def test_decompose_s():
     """Test decomposition of S gate."""
-    dev = get_custom_device_without(1, {"S", "C(S)"})
+    dev = get_custom_device_without(1, discards={"S", "C(S)"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -238,7 +256,7 @@ def decompose_s():
 
 def test_decompose_qubitunitary():
     """Test decomposition of QubitUnitary"""
-    dev = get_custom_device_without(1, {"QubitUnitary"})
+    dev = get_custom_device_without(1, discards={"QubitUnitary"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -260,7 +278,7 @@ def decompose_qubit_unitary(U: jax.core.ShapedArray([2, 2], float)):
 
 def test_decompose_singleexcitationplus():
     """Test decomposition of single excitation plus."""
-    dev = get_custom_device_without(2, {"SingleExcitationPlus", "C(SingleExcitationPlus)"})
+    dev = get_custom_device_without(2, discards={"SingleExcitationPlus", "C(SingleExcitationPlus)"})
 
     @qjit(target="mlir")
     @qml.qnode(dev)
@@ -304,3 +322,25 @@ def decompose_singleexcitationplus(theta: float):
 
 
 test_decompose_singleexcitationplus()
+
+
+def test_decompose_to_matrix():
+    """Test decomposition of QubitUnitary"""
+    dev = get_custom_device_without(1, force_matrix={"PauliY"})
+
+    @qjit(target="mlir")
+    @qml.qnode(dev)
+    # CHECK-LABEL: public @jit_decompose_to_matrix
+    def decompose_to_matrix():
+        # CHECK: quantum.custom "PauliX"
+        qml.PauliX(wires=0)
+        # CHECK: quantum.unitary
+        qml.PauliY(wires=0)
+        # CHECK: quantum.custom "PauliZ"
+        qml.PauliZ(wires=0)
+        return measure(wires=0)
+
+    print(decompose_to_matrix.mlir)
+
+
+test_decompose_to_matrix()

frontend/test/pytest/test_config_functions.py

@@ -22,7 +22,7 @@
 import pytest
 
 from catalyst.device import QJITDevice, validate_device_capabilities
-from catalyst.device.qjit_device import check_no_overlap
+from catalyst.device.qjit_device import check_no_overlap, validate_device_capabilities
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.toml import (
     DeviceCapabilities,