Enable decomposition of MultiRZ at MLIR (#2160)

**Context:**
This PR refactors the implementation of `SignatureAnalyzer` in the
decompose-lowering pass extending the analyser and the pass to support
special operator definitions (e.g., MultiRZ) along with CustomOps as new
patterns.

**Description of the Change:**
- Add a `BaseSignatureAnalyzer` class that servers a base analyzer in
decompose-lowering
- Add MultiRZ decomposition patterns to the lowering pass
- Fix the issue with trying to capture and compile unutilized abstracted
Adjoint and Controlled decomposition rules from the graph. (We don't
need these rules from the graph as we handle these symbolic ops using
Catalyst AdjointOp and ControlledOp). Avoiding this saves compilation
time and failures for templated rules.
- Include MLIR and Python tests.


**Related Issues:**
[sc-102152]

---------

Co-authored-by: Hong-Sheng Zheng <[email protected]>

Author: maliasadi

doc/releases/changelog-dev.md

@@ -4,6 +4,11 @@
 
 <h3>Improvements 🛠</h3>
 
+* The ``decompose-lowering`` MLIR pass now supports ``qml.MultiRZ``
+  with an arbitrary number of wires. This decomposition is performed
+  at MLIR when both capture and graph-decomposition are enabled.
+  [(#2160)](https://github.com/PennyLaneAI/catalyst/pull/2160)
+
 * A new option ``use_nameloc`` has been added to :func:`~.qjit` that embeds variable names
   from Python into the compiler IR, which can make it easier to read when debugging programs.
   [(#2054)](https://github.com/PennyLaneAI/catalyst/pull/2054)
@@ -21,6 +26,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes the issue with capturing unutilized abstracted adjoint and controlled rules
+  by the graph in the new decomposition framework.
+  [(#2160)](https://github.com/PennyLaneAI/catalyst/pull/2160)
+
 * Fixes the translation of plxpr control flow for edge cases where the `consts` were being
   reordered.
   [(#2128)](https://github.com/PennyLaneAI/catalyst/pull/2128)
@@ -47,6 +56,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Ali Asadi,
 Christina Lee,
 Roberto Turrado,
 Paul Haochen Wang.

frontend/catalyst/from_plxpr/decompose.py

@@ -234,7 +234,14 @@ def interpret_measurement(self, measurement: "qml.measurement.MeasurementProcess
                         num_params=num_params,
                         requires_copy=num_wires == -1,
                     )
-                else:  # pragma: no cover
+                elif not any(
+                    keyword in getattr(op.op, "name", "") for keyword in ("Adjoint", "Controlled")
+                ):  # pragma: no cover
+                    # Note that the graph-decomposition returns abstracted rules
+                    # for Adjoint and Controlled operations, so we skip them here.
+                    # These abstracted rules cannot be captured and lowered.
+                    # We use MLIR AdjointOp and ControlledOp primitives
+                    # to deal with decomposition of symbolic operations at PLxPR.
                     raise ValueError(f"Could not capture {op} without the number of wires.")
 
         data, struct = jax.tree_util.tree_flatten(measurement)

frontend/test/pytest/from_plxpr/test_from_plxpr.py

@@ -1051,6 +1051,45 @@ def circuit():
         expected = np.array([1, 0, 0, 1]) / np.sqrt(2)
         assert qml.math.allclose(circuit(), expected)
 
+        expected_resources = {"CZ": 1, "Hadamard": 3}
+        resources = qml.specs(circuit, level="device")()["resources"]
+        assert resources.gate_types == expected_resources
+
+        qml.decomposition.disable_graph()
+        qml.capture.disable()
+
+    def test_multirz(self):
+        """Test that multirz decomposition works with from_plxpr."""
+
+        qml.capture.enable()
+        qml.decomposition.enable_graph()
+
+        @partial(
+            qml.transforms.decompose,
+            gate_set={"X", "Y", "Z", "S", "H", "CNOT", "RZ", "Rot", "GlobalPhase"},
+        )
+        @qml.qnode(qml.device("lightning.qubit", wires=4))
+        def circuit():
+            qml.Hadamard(0)
+            qml.ctrl(qml.MultiRZ(0.345, wires=[1, 2]), control=0)
+            qml.adjoint(qml.MultiRZ(0.25, wires=[1, 2]))
+            qml.MultiRZ(0.5, wires=[0, 1])
+            qml.MultiRZ(0.5, wires=[0])
+            qml.MultiRZ(0.5, wires=[0, 1, 3])
+            return qml.expval(qml.X(0))
+
+        without_qjit = circuit()
+        with_qjit = qml.qjit(circuit)
+
+        assert qml.math.allclose(without_qjit, with_qjit())
+
+        # TODO: Remove this static dict when capture & graph enabled support
+        # resource counting with qml.specs via from_plxpr conversion.
+        expected_resources = {"GlobalPhase": 14, "RZ": 20, "CNOT": 22, "Hadamard": 5}
+
+        resources = qml.specs(with_qjit, level="device")()["resources"]
+        assert resources.gate_types == expected_resources
+
         qml.decomposition.disable_graph()
         qml.capture.disable()