Emit a cirq.Circuit from a squin kernel (#311)

There's still quite a few things to do here:

* Some of the operator statements (e.g. `Kron`) don't have a 1:1 mapping
to a cirq gate. They probably need a "runtime", since they are easy
enough to apply to qubits. Alternatively, we could use custom gates in
cirq, but I'm not much of a fan.
* Methods for the noise statements are still missing.
* There is some very hacky stuff I'm doing right now in order to obtain
the `BlockArguments` and `ReturnValues` in the frame for nested kernels.
There's probably a better way to do this, but I don't know how (yet).

---------

Co-authored-by: Phillip Weinberg <[email protected]>

Author: david-pl

src/bloqade/squin/cirq/__init__.py

@@ -1,12 +1,17 @@
-from typing import Any
+from typing import Any, Sequence
 
 import cirq
 from kirin import ir, types
+from kirin.emit import EmitError
 from kirin.dialects import func
 
 from . import lowering as lowering
 from .. import kernel
+
+# NOTE: just to register methods
+from .emit import op as op, qubit as qubit
 from .lowering import Squin
+from .emit.emit_circuit import EmitCirq
 
 
 def load_circuit(
@@ -87,3 +92,109 @@ def load_circuit(
         dialects=dialects,
         code=code,
     )
+
+
+def emit_circuit(
+    mt: ir.Method,
+    qubits: Sequence[cirq.Qid] | None = None,
+) -> 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:
+        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.qubit.new`
+            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.
+
+    ## Examples:
+
+    Here's a very basic example:
+
+    ```python
+    from bloqade import squin
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        h = squin.op.h()
+        squin.qubit.apply(h, q[0])
+        cx = squin.op.cx()
+        squin.qubit.apply(cx, q)
+
+    circuit = squin.cirq.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 kirin.dialects import ilist
+    from typing import Literal
+    import cirq
+
+    @squin.kernel
+    def entangle(q: ilist.IList[squin.qubit.Qubit, Literal[2]]):
+        h = squin.op.h()
+        squin.qubit.apply(h, q[0])
+        cx = squin.op.cx()
+        squin.qubit.apply(cx, q)
+        return cx
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        cx = entangle(q)
+        q2 = squin.qubit.new(3)
+        squin.qubit.apply(cx, [q[1], q2[2]])
+
+
+    # custom list of qubits on grid
+    qubits = [cirq.GridQubit(i, i+1) for i in range(5)]
+
+    circuit = squin.cirq.emit_circuit(main, 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 isinstance(mt.code, func.Function) and not mt.code.signature.output.is_subseteq(
+        types.NoneType
+    ):
+        raise EmitError(
+            "The method you are trying to convert to a circuit has a return value, but returning from a circuit is not supported."
+        )
+
+    emitter = EmitCirq(qubits=qubits)
+    return emitter.run(mt, args=())
+
+
+def dump_circuit(mt: ir.Method, qubits: Sequence[cirq.Qid] | None = None, **kwargs):
+    """Converts a squin.kernel method to a cirq.Circuit object and dumps it as JSON.
+
+    This just runs `emit_circuit` and calls the `cirq.to_json` function to emit a JSON.
+
+    Args:
+        mt (ir.Method): The kernel method from which to construct the circuit.
+
+    Keyword Args:
+        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.qubit.new`
+            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.
+
+    """
+    circuit = emit_circuit(mt, qubits=qubits)
+    return cirq.to_json(circuit, **kwargs)

src/bloqade/squin/cirq/emit/emit_circuit.py

@@ -0,0 +1,109 @@
+from typing import Sequence
+from dataclasses import field, dataclass
+
+import cirq
+from kirin import ir
+from kirin.emit import EmitABC, EmitError, EmitFrame
+from kirin.interp import MethodTable, impl
+from kirin.dialects import func
+from typing_extensions import Self
+
+from ... import kernel
+
+
+@dataclass
+class EmitCirqFrame(EmitFrame):
+    qubit_index: int = 0
+    qubits: Sequence[cirq.Qid] | None = None
+    circuit: cirq.Circuit = field(default_factory=cirq.Circuit)
+
+
+def _default_kernel():
+    return kernel
+
+
+@dataclass
+class EmitCirq(EmitABC[EmitCirqFrame, cirq.Circuit]):
+    keys = ["emit.cirq", "main"]
+    dialects: ir.DialectGroup = field(default_factory=_default_kernel)
+    void = cirq.Circuit()
+    qubits: Sequence[cirq.Qid] | None = None
+    _cached_circuit_operations: dict[int, cirq.CircuitOperation] = field(
+        init=False, default_factory=dict
+    )
+
+    def initialize(self) -> Self:
+        return super().initialize()
+
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> EmitCirqFrame:
+        return EmitCirqFrame(
+            code, has_parent_access=has_parent_access, qubits=self.qubits
+        )
+
+    def run_method(self, method: ir.Method, args: tuple[cirq.Circuit, ...]):
+        return self.run_callable(method.code, args)
+
+    def emit_block(self, frame: EmitCirqFrame, block: ir.Block) -> cirq.Circuit:
+        for stmt in block.stmts:
+            result = self.eval_stmt(frame, stmt)
+            if isinstance(result, tuple):
+                frame.set_values(stmt.results, result)
+
+        return frame.circuit
+
+
+@func.dialect.register(key="emit.cirq")
+class FuncEmit(MethodTable):
+
+    @impl(func.Function)
+    def emit_func(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Function):
+        emit.run_ssacfg_region(frame, stmt.body, ())
+        return (frame.circuit,)
+
+    @impl(func.Invoke)
+    def emit_invoke(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Invoke):
+        stmt_hash = hash((stmt.callee, stmt.inputs))
+        if (
+            cached_circuit_op := emit._cached_circuit_operations.get(stmt_hash)
+        ) is not None:
+            # NOTE: cache hit
+            frame.circuit.append(cached_circuit_op)
+            return ()
+
+        ret = stmt.result
+
+        with emit.new_frame(stmt.callee.code, has_parent_access=True) as sub_frame:
+            sub_frame.qubit_index = frame.qubit_index
+            sub_frame.qubits = frame.qubits
+
+            region = stmt.callee.callable_region
+            if len(region.blocks) > 1:
+                raise EmitError(
+                    "Subroutine with more than a single block encountered. This is not supported!"
+                )
+
+            # NOTE: get the arguments, "self" is just an empty circuit
+            method_self = emit.void
+            args = [frame.get(arg_) for arg_ in stmt.inputs]
+            emit.run_ssacfg_region(
+                sub_frame, stmt.callee.callable_region, args=(method_self, *args)
+            )
+            sub_circuit = sub_frame.circuit
+
+            # NOTE: check to see if the call terminates with a return value and fetch the value;
+            # we don't support multiple return statements via control flow so we just pick the first one
+            block = region.blocks[0]
+            return_stmt = next(
+                (stmt for stmt in block.stmts if isinstance(stmt, func.Return)), None
+            )
+            if return_stmt is not None:
+                frame.entries[ret] = sub_frame.get(return_stmt.value)
+
+        circuit_op = cirq.CircuitOperation(
+            sub_circuit.freeze(), use_repetition_ids=False
+        )
+        emit._cached_circuit_operations[stmt_hash] = circuit_op
+        frame.circuit.append(circuit_op)
+        return ()

src/bloqade/squin/cirq/emit/op.py

@@ -0,0 +1,125 @@
+import math
+
+import cirq
+import numpy as np
+from kirin.interp import MethodTable, impl
+
+from ... import op
+from .runtime import (
+    SnRuntime,
+    SpRuntime,
+    U3Runtime,
+    KronRuntime,
+    MultRuntime,
+    ScaleRuntime,
+    AdjointRuntime,
+    ControlRuntime,
+    UnitaryRuntime,
+    HermitianRuntime,
+    ProjectorRuntime,
+    OperatorRuntimeABC,
+    PauliStringRuntime,
+)
+from .emit_circuit import EmitCirq, EmitCirqFrame
+
+
+@op.dialect.register(key="emit.cirq")
+class EmitCirqOpMethods(MethodTable):
+
+    @impl(op.stmts.X)
+    @impl(op.stmts.Y)
+    @impl(op.stmts.Z)
+    @impl(op.stmts.H)
+    def hermitian(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.ConstantUnitary
+    ):
+        cirq_op = getattr(cirq, stmt.name.upper())
+        return (HermitianRuntime(cirq_op),)
+
+    @impl(op.stmts.S)
+    @impl(op.stmts.T)
+    def unitary(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.ConstantUnitary
+    ):
+        cirq_op = getattr(cirq, stmt.name.upper())
+        return (UnitaryRuntime(cirq_op),)
+
+    @impl(op.stmts.P0)
+    @impl(op.stmts.P1)
+    def projector(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.P0 | op.stmts.P1
+    ):
+        return (ProjectorRuntime(isinstance(stmt, op.stmts.P1)),)
+
+    @impl(op.stmts.Sn)
+    def sn(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Sn):
+        return (SnRuntime(),)
+
+    @impl(op.stmts.Sp)
+    def sp(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Sp):
+        return (SpRuntime(),)
+
+    @impl(op.stmts.Identity)
+    def identity(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Identity):
+        op = HermitianRuntime(cirq.IdentityGate(num_qubits=stmt.sites))
+        return (op,)
+
+    @impl(op.stmts.Control)
+    def control(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Control):
+        op: OperatorRuntimeABC = frame.get(stmt.op)
+        return (ControlRuntime(op, stmt.n_controls),)
+
+    @impl(op.stmts.Kron)
+    def kron(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Kron):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        op = KronRuntime(lhs, rhs)
+        return (op,)
+
+    @impl(op.stmts.Mult)
+    def mult(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Mult):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        op = MultRuntime(lhs, rhs)
+        return (op,)
+
+    @impl(op.stmts.Adjoint)
+    def adjoint(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Adjoint):
+        op_ = frame.get(stmt.op)
+        return (AdjointRuntime(op_),)
+
+    @impl(op.stmts.Scale)
+    def scale(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Scale):
+        op_ = frame.get(stmt.op)
+        factor = frame.get(stmt.factor)
+        return (ScaleRuntime(operator=op_, factor=factor),)
+
+    @impl(op.stmts.U3)
+    def u3(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.U3):
+        theta = frame.get(stmt.theta)
+        phi = frame.get(stmt.phi)
+        lam = frame.get(stmt.lam)
+        return (U3Runtime(theta=theta, phi=phi, lam=lam),)
+
+    @impl(op.stmts.PhaseOp)
+    def phaseop(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.PhaseOp):
+        theta = frame.get(stmt.theta)
+        op_ = HermitianRuntime(cirq.IdentityGate(num_qubits=1))
+        return (ScaleRuntime(operator=op_, factor=np.exp(1j * theta)),)
+
+    @impl(op.stmts.ShiftOp)
+    def shiftop(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.ShiftOp):
+        theta = frame.get(stmt.theta)
+
+        # NOTE: ShiftOp(theta) == U3(pi, theta, 0)
+        return (U3Runtime(math.pi, theta, 0),)
+
+    @impl(op.stmts.Reset)
+    def reset(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Reset):
+        return (HermitianRuntime(cirq.ResetChannel()),)
+
+    @impl(op.stmts.PauliString)
+    def pauli_string(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.PauliString
+    ):
+        return (PauliStringRuntime(stmt.string),)