Refactor device code around capabilities [NFC] (#1149)

Simplifies device related code around capability checks and
decomposition, and eliminates the use of `"Adjoint(*)"` and `"C(*)"`
declarations in favour of DeviceCapabilities and OperationProperties
dataclasses.

Includes the following (non-functional) changes:
- simplify `get_device_capability` signature
- simplify `catalyst_decompose` signature
- simplify `QJITDevice.__init__` signature
- eliminate `pennylane_operation_set` and use `DeviceCapabilities`
instead
- remove legacy-device `operations`, `observables`, and
`measurement_processes` properties from our device classes

Some of the code seems to have been written solely to facilitate
testing, which should be avoided.

[sc-67124]

doc/releases/changelog-dev.md

@@ -2,17 +2,17 @@
 
 <h3>New features</h3>
 
-* Experimental integration of the PennyLane capture module is available. It currently only supports 
+* Experimental integration of the PennyLane capture module is available. It currently only supports
   quantum gates, without control flow.
   [(#1109)](https://github.com/PennyLaneAI/catalyst/pull/1109)
 
   To trigger the PennyLane pipeline for capturing the program as a JaxPR, one needs to simply
   set `experimental_capture=True` in the qjit decorator.
-  
-  ```python 
+
+  ```python
   import pennylane as qml
   from catalyst import qjit
-  
+
   dev = qml.device("lightning.qubit", wires=1)
 
   @qjit(experimental_capture=True)
@@ -33,10 +33,10 @@
 
   For example,
 
-  ```python 
+  ```python
   import pennylane as qml
   from catalyst import qjit
-  
+
   dev = qml.device("lightning.qubit", wires=1, shots=((5, 2), 7))
 
   @qjit
@@ -144,8 +144,8 @@
 
 <h3>Improvements</h3>
 
-* Bufferization of `gradient.ForwardOp` and `gradient.ReverseOp` now requires 3 steps: `gradient-preprocessing`, 
-  `gradient-bufferize`, and `gradient-postprocessing`. `gradient-bufferize` has a new rewrite for `gradient.ReturnOp`. 
+* Bufferization of `gradient.ForwardOp` and `gradient.ReverseOp` now requires 3 steps: `gradient-preprocessing`,
+  `gradient-bufferize`, and `gradient-postprocessing`. `gradient-bufferize` has a new rewrite for `gradient.ReturnOp`.
   [(#1139)](https://github.com/PennyLaneAI/catalyst/pull/1139)
 
 * The decorator `self_inverses` now supports all Hermitian Gates.
@@ -167,11 +167,11 @@
 
   Three-bit Gates: Toffoli
   - [`qml.Toffoli`](https://docs.pennylane.ai/en/stable/code/api/pennylane.Toffoli.html)
-
-
 
-* Support is expanded for backend devices that exculsively return samples in the measurement 
-  basis. Pre- and post-processing now allows `qjit` to be used on these devices with `qml.expval`, 
+
+
+* Support is expanded for backend devices that exculsively return samples in the measurement
+  basis. Pre- and post-processing now allows `qjit` to be used on these devices with `qml.expval`,
   `qml.var` and `qml.probs` measurements in addiiton to `qml.sample`, using the `measurements_from_samples` transform.
   [(#1106)](https://github.com/PennyLaneAI/catalyst/pull/1106)
 
@@ -192,11 +192,11 @@
 * Update Enzyme to version `v0.0.149`.
   [(#1142)](https://github.com/PennyLaneAI/catalyst/pull/1142)
 
-* Remove the `MemMemCpyOptPass` in llvm O2 (applied for Enzyme), this reduces bugs when 
+* Remove the `MemMemCpyOptPass` in llvm O2 (applied for Enzyme), this reduces bugs when
   running gradient like functions.
   [(#1063)](https://github.com/PennyLaneAI/catalyst/pull/1063)
 
-* Functions with multiple tapes are now split with a new mlir pass `--split-multiple-tapes`, with one tape per function. 
+* Functions with multiple tapes are now split with a new mlir pass `--split-multiple-tapes`, with one tape per function.
   The reset routine that makes a maeasurement between tapes and inserts a X gate if measured one is no longer used.
   [(#1017)](https://github.com/PennyLaneAI/catalyst/pull/1017)
   [(#1130)](https://github.com/PennyLaneAI/catalyst/pull/1130)
@@ -212,6 +212,13 @@
 * The device capability loading mechanism has been moved into the `QJITDevice` constructor.
   [(#1141)](https://github.com/PennyLaneAI/catalyst/pull/1141)
 
+* Several functions related to device capabilities have been refactored.
+  [(#1149)](https://github.com/PennyLaneAI/catalyst/pull/1149)
+
+  In particular, the signatures of `get_device_capability`, `catalyst_decompose`,
+ `catalyst_acceptance`, and `QJITDevice.__init__` have changed, and the `pennylane_operation_set`
+  function has been removed entirely.
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):

frontend/catalyst/device/decomposition.py

@@ -49,10 +49,8 @@ def check_alternative_support(op, capabilities):
 
     if isinstance(op, qml.ops.Controlled):
         # "Cast" away the specialized class for gates like Toffoli, ControlledQubitUnitary, etc.
-        if (
-            capabilities.native_ops.get(op.base.name)
-            and capabilities.native_ops.get(op.base.name).controllable
-        ):
+        supported = capabilities.native_ops.get(op.base.name)
+        if supported and supported.controllable:
             return [qml.ops.Controlled(op.base, op.control_wires, op.control_values, op.work_wires)]
 
     return None
@@ -71,7 +69,7 @@ def catalyst_decomposer(op, capabilities: DeviceCapabilities):
 
     if capabilities.native_ops.get("QubitUnitary"):
         # If the device supports unitary matrices, apply the relevant conversions and fallbacks.
-        if capabilities.to_matrix_ops.get(op.name) or (
+        if op.name in capabilities.to_matrix_ops or (
             op.has_matrix and isinstance(op, qml.ops.Controlled)
         ):
             return _decompose_to_matrix(op)
@@ -81,13 +79,7 @@ def catalyst_decomposer(op, capabilities: DeviceCapabilities):
 
 @transform
 @debug_logger
-def catalyst_decompose(
-    tape: qml.tape.QuantumTape,
-    ctx,
-    stopping_condition,
-    capabilities,
-    max_expansion=None,
-):
+def catalyst_decompose(tape: qml.tape.QuantumTape, ctx, capabilities):
     """Decompose operations until the stopping condition is met.
 
     In a single call of the catalyst_decompose function, the PennyLane operations are decomposed
@@ -104,18 +96,17 @@ def catalyst_decompose(
 
     (toplevel_tape,), _ = decompose(
         tape,
-        stopping_condition,
+        stopping_condition=lambda op: bool(catalyst_acceptance(op, capabilities)),
         skip_initial_state_prep=capabilities.initial_state_prep_flag,
         decomposer=partial(catalyst_decomposer, capabilities=capabilities),
-        max_expansion=max_expansion,
         name="catalyst on this device",
         error=CompileError,
     )
 
     new_ops = []
     for op in toplevel_tape.operations:
         if has_nested_tapes(op):
-            op = _decompose_nested_tapes(op, ctx, stopping_condition, capabilities, max_expansion)
+            op = _decompose_nested_tapes(op, ctx, capabilities)
         new_ops.append(op)
     tape = qml.tape.QuantumScript(new_ops, tape.measurements, shots=tape.shots)
 
@@ -133,19 +124,15 @@ def _decompose_to_matrix(op):
     return [op]
 
 
-def _decompose_nested_tapes(op, ctx, stopping_condition, capabilities, max_expansion):
+def _decompose_nested_tapes(op, ctx, capabilities):
     new_regions = []
     for region in op.regions:
         if region.quantum_tape is None:
             new_tape = None
         else:
             with EvaluationContext.frame_tracing_context(ctx, region.trace):
                 tapes, _ = catalyst_decompose(
-                    region.quantum_tape,
-                    ctx=ctx,
-                    stopping_condition=stopping_condition,
-                    capabilities=capabilities,
-                    max_expansion=max_expansion,
+                    region.quantum_tape, ctx=ctx, capabilities=capabilities
                 )
                 new_tape = tapes[0]
         new_regions.append(
@@ -198,9 +185,22 @@ def null_postprocessing(results):
     return [new_tape], null_postprocessing
 
 
-def catalyst_acceptance(op: qml.operation.Operator, operations) -> bool:
-    """Specify whether or not an Operator is supported."""
-    return op.name in operations
+def catalyst_acceptance(op: qml.operation.Operation, capabilities: DeviceCapabilities) -> str:
+    """Check whether or not an Operator is supported."""
+    op_support = capabilities.native_ops
+
+    if isinstance(op, qml.ops.Adjoint):
+        match = catalyst_acceptance(op.base, capabilities)
+        if not match or not op_support[match].invertible:
+            return None
+    elif type(op) is qml.ops.ControlledOp:
+        match = catalyst_acceptance(op.base, capabilities)
+        if not match or not op_support[match].controllable:
+            return None
+    else:
+        match = op.name if op.name in op_support else None
+
+    return match
 
 
 @transform