DecompositionGraph can be solved with work wire constraint (#7963)

**Context:**

This is a follow-up to
#7861, updating the
decomposition graph so that it can be solved under a maximum number of
work wires constraint.

**Assumptions**

- The same number of work wires is available to every operator at the
top level of the circuit.
- All work wires will be allocated from this initial pool of wires, no
borrowing algorithmic wires (whatever that means) or reallocation of
wires not yet deallocated.
- If a decomposition rule requests work wires, these wires are kept
until the end of the decomposition. Even if the decomposition rule
deallocates these wires early in its internal implementation, since
there is currently no way to specify this, we must assume that these
wires cannot be used in further decompositions of operators within.

**Algorithm**

Before dynamic wire allocation, work wires are specified as an operator
property, and included in an operator's resource params. Two operators
with different work wire budgets (I use the term work wire budget to
mean how many work wires is available for the decomposition of an
operator) are stored as two separate nodes in the graph. With dynamic
wire allocation, the number of available work wires is no longer an
operator property, so two operators with different work wire budgets
will have identical resource representations. Therefore, we cannot use
`CompressedResourceOp` as graph nodes anymore, because it's now missing
a crucial piece of information.

The first step is to define a `_OperatorNode` class, which contains both
the `CompressedResourceOp` and the number of available work wires
(previously a part of the `CompressedResourceOp`), and use this as nodes
in the graph. The number of work wires available to each operator is
tracked during the graph construction process. Imagine passing a
`work_wires` argument to every operator at the top level of the circuit,
and each decomposition takes the wires that it needs, and pass the rest
of the work wires to every op within.

**Optimizations**

There are a few inefficiencies with this approach:

1. The graph must be constructed with a maximum work wire budget. If it
changes, the entire graph must be reconstructed with a new work wire
budget.

**Solution**: throughout the graph, we track not the number of work
wires available to an operator, but the number of work wires NOT
available to an operator. The top level operators all have
`num_work_wires_not_available=0`, indicating that they have access to
all the available work wires. If a decomposition of a top level operator
uses 2 work wires, then all operators produced from this decomposition
has `num_work_wires_not_available=2`. This allows the graph to be
constructed once and solved multiple times with different work wire
budgets.

2. The number of nodes in the graph grows significantly. Imagine every
operator now comes with a `work_wires` property. Sometimes, two
operators with different work wire budgets do not necessarily need to be
two separate nodes in the graph. For example, an `RX` with 1 work wire
available and an `RX` with 2 work wires available should not be two
separate nodes.

**Solution**: during graph construction, we track not only the number of
work wires available, but also whether any decomposition rules of an
operator, or any operator produced from any of the decompositions,
depend on work wires. Since the graph construction is recursive and
depth-first, when we first encounter an operator, and build the
sub-graph with all of its decompositions, we label this operator node as
either work-wire dependent or not. The next time we encounter the same
operator with a different work wire budget, we create a new node for
this operator only if this operator was previously marked as work-wire
dependent. This significantly reduces the number of nodes.

**Related GitHub Issues:**
[sc-94400]

---------

Co-authored-by: Isaac De Vlugt <[email protected]>
Co-authored-by: Mudit Pandey <[email protected]>

Author: 3 people

doc/releases/changelog-dev.md

@@ -255,6 +255,7 @@
 
 * With :func:`~.decomposition.enable_graph()`, dynamically allocated wires are now supported in decomposition rules. This provides a smoother overall experience when decomposing operators in a way that requires auxiliary/work wires.
   [(#7861)](https://github.com/PennyLaneAI/pennylane/pull/7861)
+  [(#7963)](https://github.com/PennyLaneAI/pennylane/pull/7963)
 
 * A :class:`~.decomposition.decomposition_graph.DecompGraphSolution` class is added to store the solution of a decomposition graph. An instance of this class is returned from the `solve` method of the :class:`~.decomposition.decomposition_graph.DecompositionGraph`.
   [(#8031)](https://github.com/PennyLaneAI/pennylane/pull/8031)

pennylane/decomposition/decomposition_graph.py

@@ -48,6 +48,11 @@ class WorkWireSpec:
     garbage: int = 0
     """Garbage wires could be allocated in any state, and can be deallocated in any state."""
 
+    @property
+    def total(self) -> int:
+        """The total number of work wires."""
+        return self.zeroed + self.borrowed + self.burnable + self.garbage
+
 
 @overload
 def register_condition(condition: Callable) -> Callable[[Callable], DecompositionRule]: ...
@@ -402,7 +407,7 @@ def is_applicable(self, *args, **kwargs) -> bool:
             return True
         return self._condition(*args, **kwargs)
 
-    def work_wire_spec(self, *args, **kwargs) -> WorkWireSpec:
+    def get_work_wire_spec(self, *args, **kwargs) -> WorkWireSpec:
         """Gets the work wire requirements of this decomposition rule"""
         if isinstance(self._work_wire_spec, dict):
             return WorkWireSpec(**self._work_wire_spec)

pennylane/decomposition/decomposition_rule.py

@@ -34,7 +34,7 @@ class Resources:
     """
 
     gate_counts: dict[CompressedResourceOp, int] = field(default_factory=dict)
-    weighted_cost: float | None = field(default=None)
+    weighted_cost: float = field(default=None)
 
     def __post_init__(self):
         """Verify that all gate counts are non-zero."""

pennylane/decomposition/resources.py

@@ -37,8 +37,9 @@ def _resource_fn(base_class, base_params):  # pylint: disable=unused-argument
             for decomp_op, count in base_resources.gate_counts.items()
         }
 
+    # pylint: disable=protected-access
     @register_condition(_condition_fn)
-    @register_resources(_resource_fn)
+    @register_resources(_resource_fn, work_wires=base_decomposition._work_wire_spec)
     def _impl(*params, wires, base, **__):
         # pylint: disable=protected-access
         qml.adjoint(base_decomposition._impl)(*params, wires=wires, **base.hyperparameters)
@@ -183,7 +184,7 @@ def decompose_to_base(*params, wires, base, **__):
 self_adjoint: DecompositionRule = decompose_to_base
 
 
-def make_controlled_decomp(base_decomposition):
+def make_controlled_decomp(base_decomposition: DecompositionRule):
     """Create a decomposition rule for the control of a decomposition rule."""
 
     def _condition_fn(base_params, **_):
@@ -209,9 +210,9 @@ def _resource_fn(
         gate_counts[resource_rep(qml.PauliX)] = num_zero_control_values * 2
         return gate_counts
 
-    # pylint: disable=too-many-arguments
+    # pylint: disable=protected-access,too-many-arguments
     @register_condition(_condition_fn)
-    @register_resources(_resource_fn)
+    @register_resources(_resource_fn, work_wires=base_decomposition._work_wire_spec)
     def _impl(*params, wires, control_wires, control_values, work_wires, work_wire_type, base, **_):
         zero_control_wires = [w for w, val in zip(control_wires, control_values) if not val]
         for w in zero_control_wires:
@@ -249,8 +250,9 @@ def _resource_fn(**resource_params):
         gate_counts[resource_rep(qml.X)] = gate_counts.get(resource_rep(qml.X), 0) + num_x * 2
         return gate_counts
 
+    # pylint: disable=protected-access
     @register_condition(_condition_fn)
-    @register_resources(_resource_fn)
+    @register_resources(_resource_fn, work_wires=inner_decomp._work_wire_spec)
     def _impl(*params, wires, control_wires, control_values, **kwargs):
         zero_control_wires = [w for w, val in zip(control_wires, control_values) if not val]
         for w in zero_control_wires:

pennylane/decomposition/symbolic_decomposition.py

@@ -147,6 +147,14 @@ def _t_ps(wires, **__):
 
 decompositions["T"] = [_t_ps]
 
+
+@qml.register_resources({qml.RZ: 3, qml.RY: 2, qml.CNOT: 2})
+def _crot(*_, **__):
+    raise NotImplementedError
+
+
+decompositions["CRot"] = [_crot]
+
 ################################################
 # Custom Decompositions For Symbolic Operators #
 ################################################

tests/decomposition/conftest.py

@@ -31,6 +31,7 @@
 )
 from pennylane.decomposition.decomposition_graph import _to_name
 from pennylane.exceptions import DecompositionError
+from pennylane.operation import Operation
 
 # pylint: disable=protected-access,no-name-in-module
 
@@ -153,7 +154,7 @@ def test_graph_construction(self, _):
     def test_graph_construction_non_applicable_rules(self, _):
         """Tests rules which are not applicable are skipped."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom op"""
 
             resource_keys = {"num_wires"}
@@ -190,7 +191,7 @@ def some_other_rule(*_, **__):
     def test_gate_set(self, _):
         """Tests that graph construction stops at the target gate set."""
 
-        class CustomOp(qml.operation.Operator):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()
@@ -247,6 +248,9 @@ def test_graph_solve(self, _):
             {qml.RY: 1, qml.GlobalPhase: 1, qml.RZ: 1},
         )
 
+        # verify that is_solved_for returns False for non-existent operators
+        assert not solution.is_solved_for(qml.Toffoli(wires=[0, 1, 2]))
+
     def test_decomposition_not_found(self, _):
         """Tests that the correct error is raised if a decomposition isn't found."""
 
@@ -258,7 +262,7 @@ def test_decomposition_not_found(self, _):
     def test_lazy_solve(self, _):
         """Tests the lazy keyword argument."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()
@@ -267,7 +271,7 @@ class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-metho
             def resource_params(self):
                 return {}
 
-        class AnotherOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class AnotherOp(Operation):  # pylint: disable=too-few-public-methods
             """Another custom operation."""
 
             resource_keys = set()
@@ -311,7 +315,7 @@ def _another_decomp(*_, **__):
     def test_decomposition_with_resource_params(self, _):
         """Tests operators with non-empty resource params."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = {"num_wires"}
@@ -358,6 +362,125 @@ def _custom_decomp(*_, **__):
             {qml.RZ: 2, qml.RX: 1, qml.GlobalPhase: 1},
         )
 
+    def test_work_wire_requirement(self, _):
+        """Tests that the graph respects the work wire requirement."""
+
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
+            """A custom operation."""
+
+            resource_keys = set()
+
+            @property
+            def resource_params(self):
+                return {}
+
+        @qml.register_resources({qml.Toffoli: 2, qml.CRot: 1}, work_wires={"zeroed": 1})
+        def _decomp_with_work_wire(*_, **__):
+            raise NotImplementedError
+
+        @qml.register_resources({qml.Toffoli: 2, qml.CRot: 3})
+        def _decomp_without_work_wire(*_, **__):
+            raise NotImplementedError
+
+        graph = DecompositionGraph(
+            [CustomOp(wires=[0, 1, 2])],
+            gate_set={qml.Toffoli, qml.CRot},
+            alt_decomps={CustomOp: [_decomp_without_work_wire, _decomp_with_work_wire]},
+        )
+
+        solution = graph.solve(num_work_wires=0)
+        assert solution.decomposition(CustomOp(wires=[0, 1, 2])) is _decomp_without_work_wire
+
+        solution = graph.solve(num_work_wires=1)
+        assert (
+            solution.decomposition(CustomOp(wires=[0, 1, 2]), num_work_wires=1)
+            is _decomp_with_work_wire
+        )
+
+    def test_multiple_nodes_with_different_work_wire_budget(self, _):
+        """Tests that the same operator produced under different work wire budgets
+        are stored as different nodes in the graph, and results can be queried."""
+
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
+            """A custom operation."""
+
+            resource_keys = set()
+
+            @property
+            def resource_params(self):
+                return {}
+
+        @qml.register_resources({qml.Toffoli: 2, qml.CRot: 1}, work_wires={"zeroed": 2})
+        def _decomp_with_work_wire(*_, **__):
+            raise NotImplementedError
+
+        @qml.register_resources({qml.Toffoli: 4, qml.CRot: 3})
+        def _decomp_without_work_wire(*_, **__):
+            raise NotImplementedError
+
+        class LargeOp(Operation):  # pylint: disable=too-few-public-methods
+            """A larger custom operation."""
+
+            resource_keys = set()
+
+            @property
+            def resource_params(self):
+                return {}
+
+        @qml.register_resources({qml.Toffoli: 2, CustomOp: 2}, work_wires={"zeroed": 1})
+        def _decomp2_with_work_wire(*_, **__):
+            raise NotImplementedError
+
+        @qml.register_resources({qml.Toffoli: 4, CustomOp: 2})
+        def _decomp2_without_work_wire(*_, **__):
+            raise NotImplementedError
+
+        op = LargeOp(wires=[0, 1, 2, 3])
+        small_op = CustomOp(wires=[0, 1, 2])
+
+        graph = DecompositionGraph(
+            [op, small_op],
+            gate_set={qml.Toffoli, qml.RZ, qml.RY, qml.CNOT},
+            alt_decomps={
+                CustomOp: [_decomp_without_work_wire, _decomp_with_work_wire],
+                LargeOp: [_decomp2_without_work_wire, _decomp2_with_work_wire],
+            },
+        )
+
+        # 1 node for LargerOp, 2 nodes for CustomOp, 1 for Toffoli, 1 for CRot, 1 for RZ,
+        # 1 for RY, 1 for CNOT, and 1 dummy starting node, 1 decomposition from CRot,
+        # node, 2 decomposition nodes from LargerOp, 2 decompositions from each CustomOp
+        assert len(graph._graph.nodes()) == 16
+        assert len(graph._graph.edges()) == 26
+
+        solution = graph.solve(num_work_wires=0)
+        assert solution.decomposition(op) is _decomp2_without_work_wire
+        assert solution.decomposition(small_op) is _decomp_without_work_wire
+
+        solution = graph.solve(num_work_wires=1)
+        assert solution.decomposition(op, num_work_wires=1) is _decomp2_with_work_wire
+        assert solution.decomposition(small_op, num_work_wires=0) is _decomp_without_work_wire
+
+        solution = graph.solve(num_work_wires=2)
+        # When there are only 2 work wires available, by construction, it is more
+        # resource efficient to use them on the CustomOp, so even where there are
+        # enough work wires to use the more efficient decomposition for the LargeOp,
+        # we should still choose the less efficient one to achieve better overall
+        # resource efficiency. Because if we use one of the work wires to decompose
+        # the LargeOp, there won't be enough work wires left to further decompose
+        # the 2 CustomOp and it would result in significantly more gates.
+        assert solution.decomposition(op, num_work_wires=2) is _decomp2_without_work_wire
+        assert solution.decomposition(small_op, num_work_wires=2) is _decomp_with_work_wire
+
+        solution = graph.solve(num_work_wires=3)
+        assert solution.decomposition(op, num_work_wires=3) is _decomp2_with_work_wire
+        assert solution.decomposition(small_op, num_work_wires=2) is _decomp_with_work_wire
+        assert solution.decomposition(small_op, num_work_wires=3) is _decomp_with_work_wire
+
+        solution = graph.solve(num_work_wires=None)
+        assert solution.decomposition(op, num_work_wires=None) is _decomp2_with_work_wire
+        assert solution.decomposition(small_op, num_work_wires=None) is _decomp_with_work_wire
+
 
 @pytest.mark.unit
 @patch(
@@ -413,7 +536,7 @@ def test_custom_controlled_op(self, _):
     def test_controlled_base_decomposition(self, _):
         """Tests applying control on the decomposition of the target operator."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()
@@ -432,7 +555,7 @@ def second_decomp(wires):
             qml.Z(wires=wires[0])
             qml.GlobalPhase(np.pi / 2, wires=wires)
 
-        class CustomControlledOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomControlledOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()
@@ -575,7 +698,7 @@ def test_adjoint_custom(self, _):
     def test_adjoint_general(self, _):
         """Tests decomposition of a generalized adjoint operation."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()
@@ -688,7 +811,7 @@ def my_adjoint_rx(theta, wires, **__):
     def test_special_pow_decomps(self, _):
         """Tests special cases for decomposing a power."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()
@@ -724,7 +847,7 @@ def resource_params(self):
     def test_general_pow_decomps(self, _):
         """Tests the more general power decomposition rules."""
 
-        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+        class CustomOp(Operation):  # pylint: disable=too-few-public-methods
             """A custom operation."""
 
             resource_keys = set()

tests/decomposition/test_decomposition_graph.py

@@ -298,7 +298,7 @@ def test_register_work_wires(self):
         def custom_decomp(*_, **__):
             raise NotImplementedError
 
-        assert custom_decomp.work_wire_spec() == WorkWireSpec(1, 3, 4, 2)
+        assert custom_decomp.get_work_wire_spec() == WorkWireSpec(1, 3, 4, 2)
 
         @register_resources(
             lambda num_wires: {qml.CNOT: num_wires},
@@ -311,4 +311,4 @@ def custom_decomp(*_, **__):
         def custom_decomp_2(*_, **__):
             raise NotImplementedError
 
-        assert custom_decomp_2.work_wire_spec(num_wires=5) == WorkWireSpec(zeroed=2, borrowed=3)
+        assert custom_decomp_2.get_work_wire_spec(num_wires=5) == WorkWireSpec(zeroed=2, borrowed=3)