Bloq counts for cirq_ft MultiControlPauli / And

[fdmalone]

The following code throws an error:

I need to double check when this stopped working, but this code worked prior to the bump in cirq_ft versions.

from functools import cached_property
from typing import Dict

import cirq
import numpy as np
from attrs import frozen
from cirq_ft.algos.multi_control_multi_target_pauli import MultiControlPauli

from qualtran import Bloq, BloqBuilder, Register, Signature, SoquetT
from qualtran.cirq_interop import CirqGateAsBloq
from qualtran.bloqs.util_bloqs import Allocate, Free, Join, Split
from qualtran.bloqs.and_bloq import And
from qualtran.resource_counting import get_bloq_counts_graph, SympySymbolAllocator
import cirq_ft

ssa = SympySymbolAllocator()
and_cv0 = ssa.new_symbol('cv0')
and_cv1 = ssa.new_symbol('cv1')

def generalize(bloq):
    """Genereralizer for THC prepare bloqs."""
    if isinstance(bloq, (Allocate, Free, Split, Join)):
        return None
    if isinstance(bloq, CirqGateAsBloq):
        if isinstance(bloq.gate, cirq_ft.algos.And) and (len(bloq.gate.cv) == 2):
            return And(cv1=and_cv0, cv2=and_cv1, adjoint=bloq.gate.adjoint)
    return bloq

@frozen
class TestAnd(Bloq):
    r"""Prepare uniform superposition state for THC.
    """
    @cached_property
    def signature(self) -> Signature:
        return Signature(
            [
                Register("a", bitsize=1),
                Register("b", bitsize=1),
                Register("c", bitsize=1),
                Register("d", bitsize=1),
            ]
        )

    def build_composite_bloq(self, bb: 'BloqBuilder', **regs: SoquetT) -> Dict[str, 'SoquetT']:
        a = regs['a']
        b = regs['b']
        c = regs['c']
        d = regs['d']
        ctrls = bb.join(np.array([a, b, c]))
        ctrls, d = bb.add(
            CirqGateAsBloq(MultiControlPauli(cvs=(1, 1, 1), target_gate=cirq.Z)),
            controls=ctrls,
            target=d,
        )
        (a, b, c) = bb.split(ctrls)
        out_regs = {'a': a, 'b': b, 'c': c, 'd': d}
        return out_regs

bloq = TestAnd()
get_bloq_counts_graph(bloq, generalizer=generalize)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
get_bloq_counts_graph(bloq, generalizer=generalize)

/projects/qualtran/qualtran/resource_counting/bloq_counts.py:177), in get_bloq_counts_graph(bloq, generalizer, ssa, keep)
    175 if bloq is None:
    176     raise ValueError("You can't generalize away the root bloq.")
--> 177 sigma = _descend_counts(bloq, g, ssa, generalizer, keep)
    178 return g, sigma

~/projects/qualtran/qualtran/resource_counting/bloq_counts.py:131), in _descend_counts(parent, g, ssa, generalizer, keep)
    128     g.add_edge(parent, child, n=n)
    130 # Do the recursive step, which will continue to mutate `g`
--> 131 child_counts = _descend_counts(child, g, ssa, generalizer, keep)
    133 # Update `sigma` with the recursion results.
    134 for k in child_counts.keys():

~/projects/qualtran/qualtran/resource_counting/bloq_counts.py:131), in _descend_counts(parent, g, ssa, generalizer, keep)
    128     g.add_edge(parent, child, n=n)
    130 # Do the recursive step, which will continue to mutate `g`
--> 131 child_counts = _descend_counts(child, g, ssa, generalizer, keep)
    133 # Update `sigma` with the recursion results.
    134 for k in child_counts.keys():

/projects/qualtran/qualtran/resource_counting/bloq_counts.py:111), in _descend_counts(parent, g, ssa, generalizer, keep)
    109     return {parent: 1}
    110 try:
--> 111     count_decomp = parent.bloq_counts(ssa)
    112 except NotImplementedError:
    113     # Base case 2: Decomposition (or `bloq_counts`) is not implemented. This is left as a
    114     #              leaf node.
    115     return {parent: 1}

~/projects/qualtran/qualtran/_infra/bloq.py:243), in Bloq.bloq_counts(self, ssa)
    234 def bloq_counts(self, ssa: Optional['SympySymbolAllocator'] = None) -> Set['BloqCountT']:
    235     """Return a set of `(n, bloq)` tuples where bloq is used `n` times in the decomposition.
    236 
    237     By default, this method will use `self.decompose_bloq()` to count up bloqs.
   (...)
    241     sympy symbols (perhaps with the aid of the provided `SympySymbolAllocator`).
    242     """
--> 243     return self.decompose_bloq().bloq_counts(ssa)

/~/projects/qualtran/qualtran/cirq_interop/_cirq_to_bloq.py:71), in CirqGateAsBloq.decompose_bloq(self)
     69 in_quregs = self.signature.get_cirq_quregs()
     70 qubit_manager = cirq.ops.SimpleQubitManager()
---> 71 cirq_op, out_quregs = self.as_cirq_op(qubit_manager, **in_quregs)
     72 context = cirq.DecompositionContext(qubit_manager=qubit_manager)
     73 decomposed_optree = cirq.decompose_once(cirq_op, context=context, default=None)

~/projects/qualtran/qualtran/cirq_interop/_cirq_to_bloq.py:148), in CirqGateAsBloq.as_cirq_op(self, qubit_manager, **cirq_quregs)
    146 if not isinstance(self.gate, cirq_ft.GateWithRegisters):
    147     return self.gate.on(*cirq_quregs['qubits'].flatten()), cirq_quregs
--> 148 return _construct_op_from_gate(
    149     self.gate,
    150     in_quregs={k: np.array(v) for k, v in cirq_quregs.items()},
    151     qubit_manager=qubit_manager,
    152 )

~/projects/qualtran/qualtran/cirq_interop/_bloq_to_cirq.py:291), in _construct_op_from_gate(gate, in_quregs, qubit_manager)
    286     all_quregs[reg.name] = np.array(qubit_manager.qalloc(reg.total_bits())).reshape(
    287         full_shape
    288     )
    289 if reg.side == cirq_ft.infra.Side.LEFT:
    290     # LEFT only registers should be de-allocated and not be part of output.
--> 291     qubit_manager.qfree(in_quregs[reg.name].flatten())
    293 if reg.side & cirq_ft.infra.Side.RIGHT:
    294     # Right registers should be part of the output.
    295     out_quregs[reg.name] = all_quregs[reg.name]

qubit_manager.py:97), in SimpleQubitManager.qfree(self, qubits)
     95 good |= isinstance(q, BorrowableQubit) and q.id < self._borrow_id
     96 if not good:
---> 97     raise ValueError(f"{q} was not allocated by {self}")

ValueError: junk[0] was not allocated by <cirq.ops.qubit_manager.SimpleQubitManager object at 0x7f62b98d3690>

[fdmalone]

Any idea @tanujkhattar ?

[tanujkhattar]

I think I know what's going on. I'll send a fix soon.

[tanujkhattar]

Explanation: The root cause of the problem here is that during the interop, we allocate the named qubits that are expected as LEFT registers by the gates using signature.get_cirq_quregs and also provide a cirq.ops.SimpleQubitManager() to manage allocations / deallocations of additional qubits. We end up with this bug when the qubit manager we provided tries to deallocate the qubits we allocated separately using signature.get_cirq_quregs.

Ideally, every qubit allocation / deallocation should be managed by the same qubit manager. So, we should either get rid of signature.get_cirq_quregs() or find a way to inform our qubit manager to manage these externally allocated qubits so they can be deallocated by the qubit manager when a request comes.

https://github.com/quantumlib/Qualtran/pull/385 adds a InteropQubitManager which does the latter. We can evolve the functionality of this InteropQubitManager in subsequent PRs so that it can more natively allocate / deallocate named qubits based on register names and shapes but for now, the fix is sufficient to unblock the interop.