Support aliased controlled gates via their base operation (#792)

Small improvements to the gate decomposition strategy with the new device API. Controlled operations defined via specialized classes (like `Toffoli` or `ControlledQubitUnitary`) are now implemented as controlled versions of their base operation if the device supports it. _In particular, `MultiControlledX` is no longer executed as a `QubitUnitary` with Lightning._ [sc-65228]
PennyLaneAI · Jun 21, 2024 · 298449f · 298449f
1 parent 4559fe6
commit 298449f
Show file tree

Hide file tree

Showing 20 changed files with 286 additions and 217 deletions.
diff --git a/.pylintrc b/.pylintrc
@@ -37,7 +37,7 @@ enable=useless-suppression
 # it should appear only once).
 # Cyclical import checks are disabled for now as they are frequently used in
 # the code base, but this can be removed in the future once cycles are resolved.
-disable=too-few-public-methods,invalid-name,too-many-locals,cyclic-import,import-error,no-else-return,unnecessary-ellipsis,duplicate-code
+disable=too-few-public-methods,invalid-name,too-many-locals,cyclic-import,import-error,no-else-return,unnecessary-ellipsis,duplicate-code,abstract-method,no-name-in-module
 
 [MISCELLANEOUS]
 

diff --git a/doc/changelog.md b/doc/changelog.md
@@ -5,7 +5,8 @@
 * `qjit` adheres to user-specified `mcm_method` given to the `QNode`.
   [(#798)](https://github.com/PennyLaneAI/catalyst/pull/798)
 
-* The `dynamic_one_shot` transform uses a single auxiliary tape which is repeatedly simulated `n_shots` times to simulate hardware-like results.
+* The `dynamic_one_shot` transform uses a single auxiliary tape which is repeatedly simulated
+  `n_shots` times to simulate hardware-like results.
   The loop over shots is executed with `catalyst.vmap`.
   [(#5617)](https://github.com/PennyLaneAI/pennylane/pull/5617)
 
@@ -91,7 +92,7 @@
 
   ```
 
-* Support for using `catalyst.value_and_grad` with a `qjit`-ted function. 
+* Support for using `catalyst.value_and_grad` with a `qjit`-ted function.
   [(#804)](https://github.com/PennyLaneAI/catalyst/pull/804)
 
   ```py
@@ -237,11 +238,16 @@
   [(#784)](https://github.com/PennyLaneAI/catalyst/pull/784)
 
 * Catalyst's adjoint and ctrl methods are now fully compatible with the PennyLane equivalent when
-  applied to a single Operator. This should lead to improved compatibility with PennyLane library code,
-  as well when reusing quantum functions with both Catalyst and PennyLane.
+  applied to a single Operator. This should lead to improved compatibility with PennyLane library
+  code, as well when reusing quantum functions with both Catalyst and PennyLane.
   [(#768)](https://github.com/PennyLaneAI/catalyst/pull/768)
   [(#771)](https://github.com/PennyLaneAI/catalyst/pull/771)
 
+* Controlled operations defined via specialized classes (like `Toffoli` or `ControlledQubitUnitary`)
+  are now implemented as controlled versions of their base operation if the device supports it.
+  In particular, `MultiControlledX` is no longer executed as a `QubitUnitary` with Lightning.
+  [(#792)](https://github.com/PennyLaneAI/catalyst/pull/792)
+
 * Catalyst now has support for `qml.sample(m)` where `m` is the result of a mid-circuit
   measurement. For now the feature is equivalent to returning `m` directly from a quantum
   function, but will be improved to return an array with one measurement result for each
@@ -280,7 +286,8 @@
   both be provided as keyword arguments.
   [(#790)](https://github.com/PennyLaneAI/catalyst/pull/790)
 
-* Finite difference is now always possible regardless of whether the differentiated function has a valid gradient for autodiff or not.
+* Finite difference is now always possible regardless of whether the differentiated function has a
+  valid gradient for autodiff or not.
   [(#789)](https://github.com/PennyLaneAI/catalyst/pull/789)
 
 * A new GitHub workflow makes available a binary distribution for Linux Arm64.
@@ -295,7 +302,9 @@
 
 <h3>Bug fixes</h3>
 
-* `device_shots` is modified to `0` on the fly in `Measure` (and set back to its original value after the call to `PartialProbs`) to compute mid-circuit probabilities analytically, even when the device has finite shots.
+* `device_shots` is modified to `0` on the fly in `Measure` (and set back to its original value
+  after the call to `PartialProbs`) to compute mid-circuit probabilities analytically, even when the
+  device has finite shots.
   [(#801)](https://github.com/PennyLaneAI/catalyst/pull/801)
 
 * The Catalyst runtime now raises an error if an qubit is accessed out of bounds from the allocated
@@ -306,7 +315,9 @@
   [(#733)](https://github.com/PennyLaneAI/catalyst/pull/733)
 
 * Correctly linking openblas routines necessary for `jax.scipy.linalg.expm`.
-  In this bug fix, four openblas routines were newly linked and are now discoverable by `stablehlo.custom_call@<blas_routine>`. They are `blas_dtrsm`, `blas_ztrsm`, `lapack_dgetrf`, `lapack_zgetrf`.
+  In this bug fix, four openblas routines were newly linked and are now discoverable by
+  `stablehlo.custom_call@<blas_routine>`. They are `blas_dtrsm`, `blas_ztrsm`, `lapack_dgetrf`,
+  `lapack_zgetrf`.
   [(#752)](https://github.com/PennyLaneAI/catalyst/pull/752)
 
 * Correctly recording types of constant array when lowering `catalyst.grad` to mlir
@@ -320,7 +331,8 @@
 
 <h3>Internal changes</h3>
 
-* Catalyst uses the `collapse` method of Lightning simulators in `Measure` to select a state vector branch and normalize.
+* Catalyst uses the `collapse` method of Lightning simulators in `Measure` to select a state vector
+  branch and normalize.
   [(#801)](https://github.com/PennyLaneAI/catalyst/pull/801)
 
 * The `QCtrl` class in Catalyst has been renamed to `HybridCtrl`, indicating its capability
@@ -413,48 +425,56 @@
   interface and allows for multiple `MemrefCallable` to be defined for a single
   callback, which is necessary for custom gradient of `pure_callbacks`.
 
-* A new `catalyst::gradient::GradientOpInterface` is available when querying the gradient method in the mlir c++ api.
+* A new `catalyst::gradient::GradientOpInterface` is available when querying the gradient method in
+  the mlir c++ api.
   [(#800)](https://github.com/PennyLaneAI/catalyst/pull/800)
 
-  `catalyst::gradient::GradOp`, `ValueAndGradOp`, `JVPOp`, and `VJPOp` now inherits traits in this new `GradientOpInterface`. The supported attributes are now `getMethod()`, `getCallee()`, `getDiffArgIndices()`, `getDiffArgIndicesAttr()`, `getFiniteDiffParam()`, and `getFiniteDiffParamAttr()`. 
-
-  - There are operations that could potentially be used as `GradOp`, `ValueAndGradOp`, `JVPOp` or `VJPOp`. When trying to get the gradient method, instead of doing 
-  ```C++
-        auto gradOp = dyn_cast<GradOp>(op);
-        auto jvpOp = dyn_cast<JVPOp>(op);
-        auto vjpOp = dyn_cast<VJPOp>(op);
-
-        llvm::StringRef MethodName;
-        if (gradOp)
-            MethodName = gradOp.getMethod();
-        else if (jvpOp)
-            MethodName = jvpOp.getMethod();
-        else if (vjpOp)
-            MethodName = vjpOp.getMethod();
-  ```
-  to identify which op it actually is and protect against segfaults (calling `nullptr.getMethod()`), in the new interface we just do 
-  ```C++
-        auto gradOpInterface = cast<GradientOpInterface>(op);
-        llvm::StringRef MethodName = gradOpInterface.getMethod();
-  ```
-
-  - Another advantage is that any concrete gradient operation object can behave like a `GradientOpInterface` :
-  ```C++
-  GradOp op; // or ValueAndGradOp op, ...
-  auto foo = [](GradientOpInterface op){
-    llvm::errs() << op.getCallee();
-  };
-  foo(op);  // this works!
-  ```
-
-  - Finally, concrete op specific methods can still be called by "reinterpret"-casting the interface back to a concrete op (provided the concrete op type is correct):
-  ```C++
-  auto foo = [](GradientOpInterface op){
-    size_t numGradients = cast<ValueAndGradOp>(&op)->getGradients().size();
-  };
-  ValueAndGradOp op;
-  foo(op);  // this works!  
-  ```
+  `catalyst::gradient::GradOp`, `ValueAndGradOp`, `JVPOp`, and `VJPOp` now inherits traits in this
+  new `GradientOpInterface`. The supported attributes are now `getMethod()`, `getCallee()`,
+  `getDiffArgIndices()`, `getDiffArgIndicesAttr()`, `getFiniteDiffParam()`, and
+  `getFiniteDiffParamAttr()`.
+
+  - There are operations that could potentially be used as `GradOp`, `ValueAndGradOp`, `JVPOp` or
+    `VJPOp`. When trying to get the gradient method, instead of doing
+    ```C++
+          auto gradOp = dyn_cast<GradOp>(op);
+          auto jvpOp = dyn_cast<JVPOp>(op);
+          auto vjpOp = dyn_cast<VJPOp>(op);
+
+          llvm::StringRef MethodName;
+          if (gradOp)
+              MethodName = gradOp.getMethod();
+          else if (jvpOp)
+              MethodName = jvpOp.getMethod();
+          else if (vjpOp)
+              MethodName = vjpOp.getMethod();
+    ```
+    to identify which op it actually is and protect against segfaults (calling
+    `nullptr.getMethod()`), in the new interface we just do
+    ```C++
+          auto gradOpInterface = cast<GradientOpInterface>(op);
+          llvm::StringRef MethodName = gradOpInterface.getMethod();
+    ```
+
+  - Another advantage is that any concrete gradient operation object can behave like a
+    `GradientOpInterface`:
+    ```C++
+    GradOp op; // or ValueAndGradOp op, ...
+    auto foo = [](GradientOpInterface op){
+      llvm::errs() << op.getCallee();
+    };
+    foo(op);  // this works!
+    ```
+
+  - Finally, concrete op specific methods can still be called by "reinterpret"-casting the interface
+    back to a concrete op (provided the concrete op type is correct):
+    ```C++
+    auto foo = [](GradientOpInterface op){
+      size_t numGradients = cast<ValueAndGradOp>(&op)->getGradients().size();
+    };
+    ValueAndGradOp op;
+    foo(op);  // this works!
+    ```
 
 <h3>Contributors</h3>
 

diff --git a/doc/dev/custom_devices.rst b/doc/dev/custom_devices.rst
@@ -262,7 +262,7 @@ headers and fields are generally required, unless stated otherwise.
         CY = { properties = [ "invertible" ] }
         CZ = { properties = [ "invertible" ] }
         PhaseShift = { properties = [ "controllable", "invertible" ] }
-        ControlledPhaseShift = { properties = [ "controllable", "invertible" ] }
+        ControlledPhaseShift = { properties = [ "invertible" ] }
         RX = { properties = [ "controllable", "invertible" ] }
         RY = { properties = [ "controllable", "invertible" ] }
         RZ = { properties = [ "controllable", "invertible" ] }
@@ -294,7 +294,7 @@ headers and fields are generally required, unless stated otherwise.
         QubitStateVector = {}
         StatePrep = {}
         ControlledQubitUnitary = {}
-        DiagonalQubitUnitary = {}
+        MultiControlledX = {}
         SingleExcitation = {}
         SingleExcitationPlus = {}
         SingleExcitationMinus = {}
@@ -310,7 +310,7 @@ headers and fields are generally required, unless stated otherwise.
         # Gates which should be translated to QubitUnitary
         [operators.gates.matrix]
 
-        MultiControlledX = {}
+        DiagonalQubitUnitary = {}
 
         # Observables supported by the device
         [operators.observables]

diff --git a/doc/dev/quick_start.rst b/doc/dev/quick_start.rst
@@ -106,12 +106,13 @@ more complex quantum circuits; see below for the list of currently supported ope
 
 .. important::
 
-   Most decomposition logic will be equivalent to PennyLane's decomposition.
-   However, decomposition logic will differ in the following cases:
+   Decomposition will generally happen in accordance with the specification provided by devices,
+   which can vary from device to device (e.g. ``default.qubit`` and ``lightning.qubit`` might
+   decompose quite differently.)
+   However, Catalyst's decomposition logic will differ in the following cases:
 
    1. All :class:`qml.Controlled <pennylane.ops.op_math.Controlled>` operations will decompose to :class:`qml.QubitUnitary <pennylane.QubitUnitary>` operations.
-   2. :class:`qml.ControlledQubitUnitary <pennylane.ControlledQubitUnitary>` operations will decompose to :class:`qml.QubitUnitary <pennylane.QubitUnitary>` operations.
-   3. The list of device-supported gates employed by Catalyst is currently different than that of the ``lightning.qubit`` device, as defined by the :class:`~.qjit_device.QJITDevice`.
+   2. The set of operations supported by Catalyst itself can in some instances lead to additional decompositions compared to the device itself.
 
 .. raw:: html
 

diff --git a/frontend/catalyst/__init__.py b/frontend/catalyst/__init__.py
@@ -39,7 +39,6 @@
 
 try:
     if INSTALLED:
-        # pylint: disable=no-name-in-module
         from catalyst._revision import __revision__  # pragma: no cover
     else:
         from subprocess import check_output

diff --git a/frontend/catalyst/compiled_functions.py b/frontend/catalyst/compiled_functions.py
@@ -36,7 +36,7 @@
     get_decomposed_signature,
     typecheck_signatures,
 )
-from catalyst.utils import wrapper  # pylint: disable=no-name-in-module
+from catalyst.utils import wrapper
 from catalyst.utils.c_template import get_template, mlir_type_to_numpy_type
 from catalyst.utils.filesystem import Directory
 from catalyst.utils.jnp_to_memref import get_ranked_memref_descriptor

diff --git a/frontend/catalyst/device/decomposition.py b/frontend/catalyst/device/decomposition.py
@@ -47,18 +47,112 @@
 logger.addHandler(logging.NullHandler())
 
 
+def check_alternative_control_support(op, capabilities):
+    """Verify that aliased controlled operations aren't supported via alternative definitions."""
+
+    if (
+        isinstance(op, qml.ControlledQubitUnitary)
+        and capabilities.native_ops.get("QubitUnitary")
+        and capabilities.native_ops.get("QubitUnitary").controllable
+    ):
+        decomp = qml.ops.Controlled(
+            qml.QubitUnitary(*op.data, wires=op.target_wires), op.control_wires, op.control_values
+        )
+    elif (
+        isinstance(op, qml.ControlledPhaseShift)
+        and capabilities.native_ops.get("PhaseShift")
+        and capabilities.native_ops.get("PhaseShift").controllable
+    ):
+        decomp = qml.ops.Controlled(
+            qml.PhaseShift(*op.data, wires=op.target_wires), op.control_wires
+        )
+    elif (
+        isinstance(op, (qml.CNOT, qml.Toffoli, qml.MultiControlledX))
+        and capabilities.native_ops.get("PauliX")
+        and capabilities.native_ops.get("PauliX").controllable
+    ):
+        decomp = qml.ops.Controlled(
+            qml.PauliX(wires=op.target_wires), op.control_wires, op.control_values, op.work_wires
+        )
+    elif (
+        isinstance(op, qml.CY)
+        and capabilities.native_ops.get("PauliY")
+        and capabilities.native_ops.get("PauliY").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.PauliY(wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, (qml.CZ, qml.CCZ))
+        and capabilities.native_ops.get("PauliZ")
+        and capabilities.native_ops.get("PauliZ").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.PauliZ(wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, qml.CRX)
+        and capabilities.native_ops.get("RX")
+        and capabilities.native_ops.get("RX").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.RX(*op.data, wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, qml.CRY)
+        and capabilities.native_ops.get("RY")
+        and capabilities.native_ops.get("RY").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.RY(*op.data, wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, qml.CRZ)
+        and capabilities.native_ops.get("RZ")
+        and capabilities.native_ops.get("RZ").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.RZ(*op.data, wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, qml.CRot)
+        and capabilities.native_ops.get("Rot")
+        and capabilities.native_ops.get("Rot").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.Rot(*op.data, wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, qml.CH)
+        and capabilities.native_ops.get("Hadamard")
+        and capabilities.native_ops.get("Hadamard").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.Hadamard(wires=op.target_wires), op.control_wires)
+    elif (
+        isinstance(op, qml.CSWAP)
+        and capabilities.native_ops.get("SWAP")
+        and capabilities.native_ops.get("SWAP").controllable
+    ):
+        decomp = qml.ops.Controlled(qml.SWAP(wires=op.target_wires), op.control_wires)
+    else:
+        decomp = None
+
+    return [decomp] if decomp else decomp
+
+
+def check_alternative_support(op, capabilities):
+    """Verify that aliased operations aren't supported via alternative definitions."""
+
+    if isinstance(op, qml.ops.Controlled):
+        return check_alternative_control_support(op, capabilities)
+
+    return None
+
+
 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.")
-    # TODO: remove hardcoded controlled to matrix decomp.
-    # Check op.has_matrix to support controlled ops without matrices:
+
+    alternative_decomp = check_alternative_support(op, capabilities)
+    if alternative_decomp is not None:
+        return alternative_decomp
+
     if capabilities.to_matrix_ops.get(op.name) or (
         op.has_matrix and isinstance(op, qml.ops.Controlled)
     ):
         return _decompose_to_matrix(op)
+
     return op.decomposition()
 
 

diff --git a/frontend/catalyst/device/qjit_device.py b/frontend/catalyst/device/qjit_device.py
@@ -85,7 +85,6 @@
     "PhaseShift",
     "PSWAP",
     "QubitUnitary",
-    "ControlledQubitUnitary",
     "Rot",
     "RX",
     "RY",
@@ -341,11 +340,9 @@ def default_expand_fn(self, circuit, max_expansion=10):
         Most decomposition logic will be equivalent to PennyLane's decomposition.
         However, decomposition logic will differ in the following cases:
 
-        1. All :class:`qml.QubitUnitary <pennylane.ops.op_math.Controlled>` operations
+        1. All unsupported :class:`qml.Controlled <pennylane.ops.op_math.Controlled>` instances
             will decompose to :class:`qml.QubitUnitary <pennylane.QubitUnitary>` operations.
-        2. :class:`qml.ControlledQubitUnitary <pennylane.ControlledQubitUnitary>` operations
-            will decompose to :class:`qml.QubitUnitary <pennylane.QubitUnitary>` operations.
-        3. The list of device-supported gates employed by Catalyst is currently different than
+        2. The list of device-supported gates employed by Catalyst is currently different than
             that of the ``lightning.qubit`` device, as defined by the
             :class:`~.qjit_device.QJITDevice`.
 

diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py
@@ -93,7 +93,7 @@
 from catalyst.utils.extra_bindings import FromElementsOp, TensorExtractOp
 from catalyst.utils.types import convert_shaped_arrays_to_tensors
 
-# pylint: disable=unused-argument,too-many-lines,too-many-statements,too-many-arguments,protected-access
+# pylint: disable=unused-argument,too-many-lines,too-many-statements,too-many-function-args,protected-access
 
 #########
 # Types #