Qubit allocation within _decompose_ is not compatible with current BFS based implementation of cirq.decompose.

[tanujkhattar]

cirq.decompose(op) currently works (roughly) as follows:

1. Call op_tree = cirq.decompose_once(op) and get the op-tree returned by calling op._decompose_().
2. For every operation in op_tree; recursively perform step-1.

Right now, when implementing _decompose_ method of an operation; we use qalloc and qfree methods to allocate / clean qubits. This pattern assumes that for all operations executed between qalloc and qfree; the allocated qubits are "in-use" and thus any qubit manager should not reuse the same qubits when recursively decomposing the operations that occur in between these two function calls.

For this assumption to be true, cirq.decompose should ideally perform a DFS on the implicit compute graph but it currently performs a BFS and as a result; before decomposing the operations that occur between qalloc and qfree function calls; the qfree call is processed and the allocated ancillas are freed. This results in a situation where only the "trivial" strategy used by SimpleQubitManager remains valid because there is no qubit reuse and thus the order in which qalloc / qfree is processed does not matter. But, for every other qubit allocation strategy that makes any reuse; the current implementation of cirq.decompose can result in reusing "occupied" qubits and thus lead to errors.

A simple example is shown below.

from attrs import frozen

@frozen
class PhaseWithDecompose(cirq.Gate):
    recurse: bool = True

    def _num_qubits_(self):
        return 1

    def _decompose_(self, q):
        anc = cq.qalloc(1)
        yield cirq.CX(*q, *anc)
        yield PhaseWithDecompose(recurse=False).on(*anc) if self.recurse else cirq.Z(*anc)
        yield cirq.CX(*q, *anc)
        cq.qfree(anc)

Decomposing the above gate with a simple qubit manager strategy outputs the correct expected circuit:

q = cirq.NamedQubit("phase_target")
with cq.memory_management_context():
    print(cirq.Circuit(cirq.decompose(PhaseWithDecompose().on(q))))
    print(compute_unitary(PhaseWithDecompose().on(q)))

_c0: ────────────Y^-0.5───@───Y^0.5───@────────────────────────@───Y^-0.5───@───Y^0.5───
                          │           │                        │            │
_c1: ────────────Y^-0.5───┼───────────@───Y^0.5───Z───Y^-0.5───@───Y^0.5────┼───────────
                          │                                                 │
phase_target: ────────────@─────────────────────────────────────────────────@───────────
[[ 1.+0.j  0.+0.j]
 [ 0.+0.j -1.+0.j]]

But calling it with a cq.GreedyQubitManager('ancilla', maximize_reuse=True) results in a qubit reuse error:

with cq.memory_management_context(cq.GreedyQubitManager('ancilla', maximize_reuse=True)):
    print(cirq.Circuit(cirq.decompose(PhaseWithDecompose().on(q))))

ValueError: Duplicate qids for <cirq.CNOT>. Expected unique qids but got <[cirq.NamedQubit('ancilla_0'), cirq.NamedQubit('ancilla_0')]>.

cc @NoureldinYosri This is the bug I was talking about in the last resource estimation sync.

[NoureldinYosri]

@tanujkhattar is the goal is to fix cirq.decompose or to write a new one for cirq-qubitization?

[tanujkhattar]

I think we should also consider inserting dummy operations representing allocation / deallocation requests within the circuit. This is how Bloqs does it and I think it'll be future proof to not depend on any specific implementation of tree traversal.

So, within _decompose_with_qubit_manager_, the use of qalloc qfree would look something like this?

def _decompose_with_qubit_manager_(self, qubits, qm: 'QubitManager') -> cirq.OP_TREE:
    ancilla = qm.qalloc(10) # It would be nice if we can also yield here, but I think it'll end up making the this a 2 line statement, which is annoying. Maybe we could "automatically" insert an "Allocation" gate for every new qubit which is part of the returned `cirq.OP_TREE` ?
    ...
    ...
    yield qm.qfree(ancilla) # This explicitly marks the point where `ancilla` is deallocated; which will be important for a qubit manager. 

I'm open to suggestions here!

[tanujkhattar]

Another option would be to restrict the use of only SimpleQubitManager as part of "constructing the compute graph" within _decompose_; which avoids any such issues and then make other utilities like 1) Simulation, 2) Computing the unitary and 3) Performing post processing via qubit allocation transformer better / more scalable so that using a dumb strategy like SimpleQubitManager to construct the circuit doesn't lead to any major problems.

This was my original idea because when I initially proposed the design of qalloc / qfree and I think this aligns with the philosophy that "we should not mix qubit allocation strategies with construction of the compute graph"

[NoureldinYosri]

+1 for the SimpleQubitManager restriction

replacing the BFS with DFS order in cirq.decompose should be easy-ish unless there are some weired tests that depend on the order or someone gets broken by it.

[tanujkhattar]

First step towards fixing this: quantumlib/Cirq#6116

[tanujkhattar]

Second PR which should fix this issue. See the test added using mocks as part of the PR: quantumlib/Cirq#6117

[tanujkhattar]

This is now (theoretically) fixed. Can be closed until we see a bug or exception pop up.