bloqade-circuit 0.6.4__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/cirq_utils/emit/base.py

@@ -0,0 +1,246 @@
+from typing import Sequence
+from warnings import warn
+from dataclasses import field, dataclass
+
+import cirq
+from kirin import ir, types, interp
+from kirin.emit import EmitABC, EmitFrame
+from kirin.interp import MethodTable, impl
+from kirin.dialects import py, func, ilist
+from typing_extensions import Self
+
+from bloqade.squin import kernel
+from bloqade.rewrite.passes import AggressiveUnroll
+
+
+def emit_circuit(
+    mt: ir.Method,
+    qubits: Sequence[cirq.Qid] | None = None,
+    circuit_qubits: Sequence[cirq.Qid] | None = None,
+    args: tuple = (),
+    ignore_returns: bool = False,
+) -> cirq.Circuit:
+    """Converts a squin.kernel method to a cirq.Circuit object.
+
+    Args:
+        mt (ir.Method): The kernel method from which to construct the circuit.
+
+    Keyword Args:
+        circuit_qubits (Sequence[cirq.Qid] | None):
+            A list of qubits to use as the qubits in the circuit. Defaults to None.
+            If this is None, then `cirq.LineQubit`s are inserted for every `squin.qalloc`
+            statement in the order they appear inside the kernel.
+            **Note**: If a list of qubits is provided, make sure that there is a sufficient
+            number of qubits for the resulting circuit.
+        args (tuple):
+            The arguments of the kernel function from which to emit a circuit.
+        ignore_returns (bool):
+            If `False`, emitting a circuit from a kernel that returns a value will error.
+            Set it to `True` in order to ignore the return value(s). Defaults to `False`.
+
+    ## Examples:
+
+    Here's a very basic example:
+
+    ```python
+    from bloqade import squin
+    from bloqade.cirq_utils import emit_circuit
+
+    @squin.kernel
+    def main():
+        q = squin.qalloc(2)
+        squin.h(q[0])
+        squin.cx(q[0], q[1])
+
+    circuit = emit_circuit(main)
+
+    print(circuit)
+    ```
+
+    You can also compose multiple kernels. Those are emitted as subcircuits within the "main" circuit.
+    Subkernels can accept arguments and return a value.
+
+    ```python
+    from bloqade import squin
+    from bloqade.cirq_utils import emit_circuit
+    from kirin.dialects import ilist
+    from typing import Literal
+    import cirq
+
+    @squin.kernel
+    def entangle(q: ilist.IList[squin.qubit.Qubit, Literal[2]]):
+        squin.h(q[0])
+        squin.cx(q[0], q[1])
+
+    @squin.kernel
+    def main():
+        q = squin.qalloc(2)
+        q2 = squin.qalloc(3)
+        squin.cx(q[1], q2[2])
+
+
+    # custom list of qubits on grid
+    qubits = [cirq.GridQubit(i, i+1) for i in range(5)]
+
+    circuit = emit_circuit(main, circuit_qubits=qubits)
+    print(circuit)
+
+    ```
+
+    We also passed in a custom list of qubits above. This allows you to provide a custom geometry
+    and manipulate the qubits in other circuits directly written in cirq as well.
+    """
+
+    if circuit_qubits is None and qubits is not None:
+        circuit_qubits = qubits
+        warn(
+            "The keyword argument `qubits` is deprecated. Use `circuit_qubits` instead."
+        )
+
+    if (
+        not ignore_returns
+        and isinstance(mt.code, func.Function)
+        and not mt.code.signature.output.is_subseteq(types.NoneType)
+    ):
+        raise interp.exceptions.InterpreterError(
+            "The method you are trying to convert to a circuit has a return value, but returning from a circuit is not supported."
+            " Set `ignore_returns = True` in order to simply ignore the return values and emit a circuit."
+        )
+
+    if len(args) != len(mt.args):
+        raise ValueError(
+            f"The method from which you're trying to emit a circuit takes {len(mt.args)} as input, but you passed in {len(args)} via the `args` keyword!"
+        )
+
+    emitter = EmitCirq(qubits=circuit_qubits)
+
+    symbol_op_trait = mt.code.get_trait(ir.SymbolOpInterface)
+    if (symbol_op_trait := mt.code.get_trait(ir.SymbolOpInterface)) is None:
+        raise interp.exceptions.InterpreterError(
+            "The method is not a symbol, cannot emit circuit!"
+        )
+
+    sym_name = symbol_op_trait.get_sym_name(mt.code).unwrap()
+
+    if (signature_trait := mt.code.get_trait(ir.HasSignature)) is None:
+        raise interp.exceptions.InterpreterError(
+            f"The method {sym_name} does not have a signature, cannot emit circuit!"
+        )
+
+    signature = signature_trait.get_signature(mt.code)
+    new_signature = func.Signature(inputs=(), output=signature.output)
+
+    callable_region = mt.callable_region.clone()
+    entry_block = callable_region.blocks[0]
+    args_ssa = list(entry_block.args)
+    first_stmt = entry_block.first_stmt
+
+    assert first_stmt is not None, "Method has no statements!"
+    if len(args_ssa) - 1 != len(args):
+        raise interp.exceptions.InterpreterError(
+            f"The method {sym_name} takes {len(args_ssa) - 1} arguments, but you passed in {len(args)} via the `args` keyword!"
+        )
+
+    for arg, arg_ssa in zip(args, args_ssa[1:], strict=True):
+        (value := py.Constant(arg)).insert_before(first_stmt)
+        arg_ssa.replace_by(value.result)
+        entry_block.args.delete(arg_ssa)
+
+    new_func = func.Function(
+        sym_name=sym_name, body=callable_region, signature=new_signature
+    )
+    mt_ = ir.Method(
+        dialects=mt.dialects,
+        code=new_func,
+        sym_name=sym_name,
+    )
+
+    AggressiveUnroll(mt_.dialects).fixpoint(mt_)
+    emitter.initialize()
+    emitter.run(mt_)
+    return emitter.circuit
+
+
+@dataclass
+class EmitCirqFrame(EmitFrame):
+    qubit_index: int = 0
+    qubits: Sequence[cirq.Qid] | None = None
+
+
+def _default_kernel():
+    return kernel
+
+
+@dataclass
+class EmitCirq(EmitABC[EmitCirqFrame, cirq.Circuit]):
+    keys = ("emit.cirq", "emit.main")
+    dialects: ir.DialectGroup = field(default_factory=_default_kernel)
+    void = cirq.Circuit()
+    qubits: Sequence[cirq.Qid] | None = None
+    circuit: cirq.Circuit = field(default_factory=cirq.Circuit)
+
+    def initialize(self) -> Self:
+        return super().initialize()
+
+    def initialize_frame(
+        self, node: ir.Statement, *, has_parent_access: bool = False
+    ) -> EmitCirqFrame:
+        return EmitCirqFrame(
+            node, has_parent_access=has_parent_access, qubits=self.qubits
+        )
+
+    def reset(self):
+        self.circuit = cirq.Circuit()
+
+    def eval_fallback(self, frame: EmitCirqFrame, node: ir.Statement) -> tuple:
+        return tuple(None for _ in range(len(node.results)))
+
+
+@func.dialect.register(key="emit.cirq")
+class __FuncEmit(MethodTable):
+
+    @impl(func.Function)
+    def emit_func(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Function):
+        for block in stmt.body.blocks:
+            frame.current_block = block
+            for s in block.stmts:
+                frame.current_stmt = s
+                stmt_results = emit.frame_eval(frame, s)
+                if isinstance(stmt_results, tuple):
+                    if len(stmt_results) != 0:
+                        frame.set_values(s.results, stmt_results)
+                    continue
+
+        return (emit.circuit,)
+
+    @impl(func.Invoke)
+    def emit_invoke(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Invoke):
+        raise interp.exceptions.InterpreterError(
+            "Function invokes should need to be inlined! "
+            "If you called the emit_circuit method, that should have happened, please report this issue."
+        )
+
+
+@py.indexing.dialect.register(key="emit.cirq")
+class __Concrete(interp.MethodTable):
+
+    @interp.impl(py.indexing.GetItem)
+    def getindex(self, interp, frame: interp.Frame, stmt: py.indexing.GetItem):
+        # NOTE: no support for indexing into single statements in cirq
+        return ()
+
+    @interp.impl(py.Constant)
+    def emit_constant(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: py.Constant):
+        return (stmt.value.data,)  # pyright: ignore[reportAttributeAccessIssue]
+
+
+@ilist.dialect.register(key="emit.cirq")
+class __IList(interp.MethodTable):
+    @interp.impl(ilist.New)
+    def new_ilist(
+        self,
+        emit: EmitCirq,
+        frame: interp.Frame,
+        stmt: ilist.New,
+    ):
+        return (ilist.IList(data=frame.get_values(stmt.values)),)

bloqade/cirq_utils/emit/gate.py

@@ -0,0 +1,104 @@
+import math
+
+import cirq
+from kirin.interp import MethodTable, impl
+
+from bloqade.squin import gate
+
+from .base import EmitCirq, EmitCirqFrame
+
+
+@gate.dialect.register(key="emit.cirq")
+class __EmitCirqGateMethods(MethodTable):
+
+    @impl(gate.stmts.X)
+    @impl(gate.stmts.Y)
+    @impl(gate.stmts.Z)
+    @impl(gate.stmts.H)
+    def hermitian(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.SingleQubitGate
+    ):
+        qubits = frame.get(stmt.qubits)
+        cirq_op = getattr(cirq, stmt.name.upper())
+        emit.circuit.append(cirq_op.on_each(qubits))
+        return ()
+
+    @impl(gate.stmts.S)
+    @impl(gate.stmts.T)
+    def unitary(
+        self,
+        emit: EmitCirq,
+        frame: EmitCirqFrame,
+        stmt: gate.stmts.SingleQubitNonHermitianGate,
+    ):
+        qubits = frame.get(stmt.qubits)
+        cirq_op = getattr(cirq, stmt.name.upper())
+        if stmt.adjoint:
+            cirq_op = cirq_op ** (-1)
+
+        emit.circuit.append(cirq_op.on_each(qubits))
+        return ()
+
+    @impl(gate.stmts.SqrtX)
+    @impl(gate.stmts.SqrtY)
+    def sqrt(
+        self,
+        emit: EmitCirq,
+        frame: EmitCirqFrame,
+        stmt: gate.stmts.SqrtX | gate.stmts.SqrtY,
+    ):
+        qubits = frame.get(stmt.qubits)
+
+        exponent = 0.5
+        if stmt.adjoint:
+            exponent *= -1
+
+        if isinstance(stmt, gate.stmts.SqrtX):
+            cirq_op = cirq.XPowGate(exponent=exponent)
+        else:
+            cirq_op = cirq.YPowGate(exponent=exponent)
+
+        emit.circuit.append(cirq_op.on_each(qubits))
+        return ()
+
+    @impl(gate.stmts.CX)
+    @impl(gate.stmts.CY)
+    @impl(gate.stmts.CZ)
+    def control(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.ControlledGate
+    ):
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        cirq_op = getattr(cirq, stmt.name.upper())
+        cirq_qubits = [(ctrl, target) for ctrl, target in zip(controls, targets)]
+        emit.circuit.append(cirq_op.on_each(cirq_qubits))
+        return ()
+
+    @impl(gate.stmts.Rx)
+    @impl(gate.stmts.Ry)
+    @impl(gate.stmts.Rz)
+    def rot(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.RotationGate):
+        qubits = frame.get(stmt.qubits)
+
+        turns = frame.get(stmt.angle)
+        angle = turns * 2 * math.pi
+        cirq_op = getattr(cirq, stmt.name.title())(rads=angle)
+
+        emit.circuit.append(cirq_op.on_each(qubits))
+        return ()
+
+    @impl(gate.stmts.U3)
+    def u3(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.U3):
+        qubits = frame.get(stmt.qubits)
+
+        theta = frame.get(stmt.theta) * 2 * math.pi
+        phi = frame.get(stmt.phi) * 2 * math.pi
+        lam = frame.get(stmt.lam) * 2 * math.pi
+
+        emit.circuit.append(cirq.Rz(rads=lam).on_each(*qubits))
+
+        emit.circuit.append(cirq.Ry(rads=theta).on_each(*qubits))
+
+        emit.circuit.append(cirq.Rz(rads=phi).on_each(*qubits))
+
+        return ()

bloqade/cirq_utils/emit/noise.py

@@ -0,0 +1,90 @@
+import cirq
+from kirin.interp import MethodTable, impl
+
+from bloqade.squin import noise
+
+from .base import EmitCirq, EmitCirqFrame
+
+
+@noise.dialect.register(key="emit.cirq")
+class __EmitCirqNoiseMethods(MethodTable):
+
+    two_qubit_paulis = (
+        "IX",
+        "IY",
+        "IZ",
+        "XI",
+        "XX",
+        "XY",
+        "XZ",
+        "YI",
+        "YX",
+        "YY",
+        "YZ",
+        "ZI",
+        "ZX",
+        "ZY",
+        "ZZ",
+    )
+
+    @impl(noise.stmts.Depolarize)
+    def depolarize(
+        self, interp: EmitCirq, frame: EmitCirqFrame, stmt: noise.stmts.Depolarize
+    ):
+        p = frame.get(stmt.p)
+        qubits = frame.get(stmt.qubits)
+        cirfq_op = cirq.depolarize(p, n_qubits=1).on_each(qubits)
+        interp.circuit.append(cirfq_op)
+        return ()
+
+    @impl(noise.stmts.Depolarize2)
+    def depolarize2(
+        self, interp: EmitCirq, frame: EmitCirqFrame, stmt: noise.stmts.Depolarize2
+    ):
+        p = frame.get(stmt.p)
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        cirq_qubits = [(ctrl, target) for ctrl, target in zip(controls, targets)]
+        cirq_op = cirq.depolarize(p, n_qubits=2).on_each(cirq_qubits)
+        interp.circuit.append(cirq_op)
+        return ()
+
+    @impl(noise.stmts.SingleQubitPauliChannel)
+    def single_qubit_pauli_channel(
+        self,
+        interp: EmitCirq,
+        frame: EmitCirqFrame,
+        stmt: noise.stmts.SingleQubitPauliChannel,
+    ):
+        px = frame.get(stmt.px)
+        py = frame.get(stmt.py)
+        pz = frame.get(stmt.pz)
+        qubits = frame.get(stmt.qubits)
+
+        cirq_op = cirq.asymmetric_depolarize(px, py, pz).on_each(qubits)
+        interp.circuit.append(cirq_op)
+
+        return ()
+
+    @impl(noise.stmts.TwoQubitPauliChannel)
+    def two_qubit_pauli_channel(
+        self,
+        interp: EmitCirq,
+        frame: EmitCirqFrame,
+        stmt: noise.stmts.TwoQubitPauliChannel,
+    ):
+        ps = frame.get(stmt.probabilities)
+        error_probabilities = {
+            key: p for (key, p) in zip(self.two_qubit_paulis, ps) if p != 0
+        }
+
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        cirq_qubits = [(ctrl, target) for ctrl, target in zip(controls, targets)]
+
+        cirq_op = cirq.asymmetric_depolarize(
+            error_probabilities=error_probabilities
+        ).on_each(cirq_qubits)
+        interp.circuit.append(cirq_op)
+
+        return ()

bloqade/cirq_utils/emit/qubit.py

@@ -0,0 +1,35 @@
+import cirq
+from kirin.interp import MethodTable, impl
+
+from bloqade.qubit import stmts as qubit
+
+from .base import EmitCirq, EmitCirqFrame
+
+
+@qubit.dialect.register(key="emit.cirq")
+class EmitCirqQubitMethods(MethodTable):
+    @impl(qubit.New)
+    def new(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.New):
+        if frame.qubits is not None:
+            cirq_qubit = frame.qubits[frame.qubit_index]
+        else:
+            cirq_qubit = cirq.LineQubit(frame.qubit_index)
+
+        frame.qubit_index += 1
+        return (cirq_qubit,)
+
+    @impl(qubit.Measure)
+    def measure_qubit_list(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Measure
+    ):
+        qbits = frame.get(stmt.qubits)
+        emit.circuit.append(cirq.measure(qbits))
+        return (emit.void,)
+
+    @impl(qubit.Reset)
+    def reset(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Reset):
+        qubits = frame.get(stmt.qubits)
+        emit.circuit.append(
+            cirq.ResetChannel().on_each(*qubits),
+        )
+        return ()