bloqade-circuit 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

bloqade/qasm2/parse/lowering.py

@@ -2,9 +2,9 @@ from typing import Any
 from dataclasses import field, dataclass
 
 from kirin import ir, types, lowering
-from kirin.dialects import cf, func, ilist
+from kirin.dialects import cf, scf, func, ilist
 
-from bloqade.qasm2.types import CRegType, QRegType
+from bloqade.qasm2.types import CRegType, QRegType, QubitType
 from bloqade.qasm2.dialects import uop, core, expr, glob, noise, parallel
 
 from . import ast
@@ -101,6 +101,13 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     def lower_global(
         self, state: lowering.State[ast.Node], node: ast.Node
     ) -> lowering.LoweringABC.Result:
+        if isinstance(node, ast.Name):
+            # NOTE: might be a lookup for a gate function invoke
+            try:
+                return lowering.LoweringABC.Result(state.current_frame.globals[node.id])
+            except KeyError:
+                pass
+
         raise lowering.BuildError("Global variables are not supported in QASM 2.0")
 
     def visit_MainProgram(self, state: lowering.State[ast.Node], node: ast.MainProgram):
@@ -171,7 +178,6 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     def visit_Reset(self, state: lowering.State[ast.Node], node: ast.Reset):
         state.current_frame.push(core.Reset(qarg=state.lower(node.qarg).expect_one()))
 
-    # TODO: clean this up? copied from cf dialect with a small modification
     def visit_IfStmt(self, state: lowering.State[ast.Node], node: ast.IfStmt):
         cond_stmt = core.CRegEq(
             lhs=state.lower(node.cond.lhs).expect_one(),
@@ -179,84 +185,23 @@ class QASM2(lowering.LoweringABC[ast.Node]):
         )
         cond = state.current_frame.push(cond_stmt).result
         frame = state.current_frame
-        before_block = frame.curr_block
 
-        with state.frame(node.body, region=frame.curr_region) as if_frame:
+        with state.frame(node.body) as if_frame:
             true_cond = if_frame.entr_block.args.append_from(types.Bool, cond.name)
             if cond.name:
                 if_frame.defs[cond.name] = true_cond
 
+            # NOTE: pass in definitions from outer scope (usually just for the qreg)
+            if_frame.defs.update(frame.defs)
+
             if_frame.exhaust()
-            self.branch_next_if_not_terminated(if_frame)
 
-        with state.frame([], region=frame.curr_region) as else_frame:
-            true_cond = else_frame.entr_block.args.append_from(types.Bool, cond.name)
-            if cond.name:
-                else_frame.defs[cond.name] = true_cond
-            else_frame.exhaust()
-            self.branch_next_if_not_terminated(else_frame)
-
-        with state.frame(frame.stream.split(), region=frame.curr_region) as after_frame:
-            after_frame.defs.update(frame.defs)
-            phi: set[str] = set()
-            for name in if_frame.defs.keys():
-                if frame.get(name):
-                    phi.add(name)
-                elif name in else_frame.defs:
-                    phi.add(name)
-
-            for name in else_frame.defs.keys():
-                if frame.get(name):  # not defined in if_frame
-                    phi.add(name)
-
-            for name in phi:
-                after_frame.defs[name] = after_frame.entr_block.args.append_from(
-                    types.Any, name
-                )
+            # NOTE: qasm2 can never yield anything from if
+            if_frame.push(scf.Yield())
 
-            after_frame.exhaust()
-            self.branch_next_if_not_terminated(after_frame)
-            after_frame.next_block.stmts.append(
-                cf.Branch(arguments=(), successor=frame.next_block)
-            )
+            then_body = if_frame.curr_region
 
-        if_args = []
-        for name in phi:
-            if value := if_frame.get(name):
-                if_args.append(value)
-            else:
-                raise lowering.BuildError(f"undefined variable {name} in if branch")
-
-        else_args = []
-        for name in phi:
-            if value := else_frame.get(name):
-                else_args.append(value)
-            else:
-                raise lowering.BuildError(f"undefined variable {name} in else branch")
-
-        if_frame.next_block.stmts.append(
-            cf.Branch(
-                arguments=tuple(if_args),
-                successor=after_frame.entr_block,
-            )
-        )
-        else_frame.next_block.stmts.append(
-            cf.Branch(
-                arguments=tuple(else_args),
-                successor=after_frame.entr_block,
-            )
-        )
-        before_block.stmts.append(
-            cf.ConditionalBranch(
-                cond=cond,
-                then_arguments=(cond,),
-                then_successor=if_frame.entr_block,
-                else_arguments=(cond,),
-                else_successor=else_frame.entr_block,
-            )
-        )
-        frame.defs.update(after_frame.defs)
-        frame.jump_next_block()
+        state.current_frame.push(scf.IfElse(cond, then_body=then_body))
 
     def branch_next_if_not_terminated(self, frame: lowering.Frame):
         """Branch to the next block if the current block is not terminated.
@@ -430,7 +375,56 @@ class QASM2(lowering.LoweringABC[ast.Node]):
             raise lowering.BuildError(f"Include {node.filename} not found")
 
     def visit_Gate(self, state: lowering.State[ast.Node], node: ast.Gate):
-        raise NotImplementedError("Gate lowering not supported")
+        arg_names = node.cparams + node.qparams
+        arg_types = [types.Float for _ in node.cparams] + [
+            QubitType for _ in node.qparams
+        ]
+
+        with state.frame(
+            stmts=node.body,
+            finalize_next=False,
+        ) as body_frame:
+            # NOTE: insert _self as arg
+            body_frame.curr_block.args.append_from(
+                types.Generic(
+                    ir.Method, types.Tuple.where(tuple(arg_types)), types.NoneType
+                ),
+                name=node.name + "_self",
+            )
+
+            for arg_type, arg_name in zip(arg_types, arg_names):
+                # NOTE: append args as block arguments
+                block_arg = body_frame.curr_block.args.append_from(
+                    arg_type, name=arg_name
+                )
+
+                # NOTE: add arguments as definitions to frame
+                body_frame.defs[arg_name] = block_arg
+
+            body_frame.exhaust()
+
+            # NOTE: append none as return value
+            return_val = func.ConstantNone()
+            body_frame.push(return_val)
+            body_frame.push(func.Return(return_val))
+
+            body = body_frame.curr_region
+
+        gate_func = expr.GateFunction(
+            sym_name=node.name,
+            signature=func.Signature(inputs=tuple(arg_types), output=types.NoneType),
+            body=body,
+        )
+
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=node.name,
+            dialects=self.dialects,
+            arg_names=[*node.cparams, *node.qparams],
+            code=gate_func,
+        )
+        state.current_frame.globals[node.name] = mt
 
     def visit_Instruction(self, state: lowering.State[ast.Node], node: ast.Instruction):
         params = [state.lower(param).expect_one() for param in node.params]

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)

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 ()

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),)

bloqade/squin/cirq/emit/qubit.py

@@ -0,0 +1,60 @@
+import cirq
+from kirin.interp import MethodTable, impl
+
+from ... import qubit
+from .op import OperatorRuntimeABC
+from .emit_circuit 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):
+        n_qubits = frame.get(stmt.n_qubits)
+
+        if frame.qubits is not None:
+            cirq_qubits = [frame.qubits[i + frame.qubit_index] for i in range(n_qubits)]
+        else:
+            cirq_qubits = [
+                cirq.LineQubit(i + frame.qubit_index) for i in range(n_qubits)
+            ]
+
+        frame.qubit_index += n_qubits
+        return (cirq_qubits,)
+
+    @impl(qubit.Apply)
+    def apply(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Apply):
+        op: OperatorRuntimeABC = frame.get(stmt.operator)
+        qbits = frame.get(stmt.qubits)
+        operations = op.apply(qbits)
+        for operation in operations:
+            frame.circuit.append(operation)
+        return ()
+
+    @impl(qubit.Broadcast)
+    def broadcast(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Broadcast):
+        op = frame.get(stmt.operator)
+        qbits = frame.get(stmt.qubits)
+
+        cirq_ops = []
+        for qbit in qbits:
+            cirq_ops.extend(op.apply([qbit]))
+
+        frame.circuit.append(cirq.Moment(cirq_ops))
+        return ()
+
+    @impl(qubit.MeasureQubit)
+    def measure_qubit(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.MeasureQubit
+    ):
+        qbit = frame.get(stmt.qubit)
+        frame.circuit.append(cirq.measure(qbit))
+        return ()
+
+    @impl(qubit.MeasureQubitList)
+    def measure_qubit_list(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.MeasureQubitList
+    ):
+        qbits = frame.get(stmt.qubits)
+        frame.circuit.append(cirq.measure(qbits))
+        return ()

bloqade/squin/cirq/emit/runtime.py

@@ -0,0 +1,234 @@
+import math
+from typing import Sequence
+from numbers import Number
+from dataclasses import dataclass
+
+import cirq
+
+
+@dataclass
+class OperatorRuntimeABC:
+    def num_qubits(self) -> int: ...
+
+    def check_qubits(self, qubits: Sequence[cirq.Qid]):
+        assert self.num_qubits() == len(qubits)
+
+    def apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        self.check_qubits(qubits)
+        return self.unsafe_apply(qubits, adjoint=adjoint)
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]: ...
+
+
+@dataclass
+class UnsafeOperatorRuntimeABC(OperatorRuntimeABC):
+    def check_qubits(self, qubits: Sequence[cirq.Qid]):
+        # NOTE: let's let cirq check this one
+        pass
+
+
+@dataclass
+class BasicOpRuntime(UnsafeOperatorRuntimeABC):
+    gate: cirq.Gate
+
+    def num_qubits(self) -> int:
+        return self.gate.num_qubits()
+
+
+@dataclass
+class UnitaryRuntime(BasicOpRuntime):
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        exponent = (-1) ** adjoint
+        return [self.gate(*qubits) ** exponent]
+
+
+@dataclass
+class HermitianRuntime(BasicOpRuntime):
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        return [self.gate(*qubits)]
+
+
+@dataclass
+class ProjectorRuntime(UnsafeOperatorRuntimeABC):
+    target_state: bool
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        # NOTE: this doesn't scale well, but works
+        sign = (-1) ** self.target_state
+        p = (1 + sign * cirq.Z(*qubits)) / 2
+        return [p]
+
+
+@dataclass
+class SpRuntime(UnsafeOperatorRuntimeABC):
+    def num_qubits(self) -> int:
+        return 1
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        if adjoint:
+            return SnRuntime().unsafe_apply(qubits, adjoint=False)
+
+        return [(cirq.X(*qubits) - 1j * cirq.Y(*qubits)) / 2]  # type: ignore  -- we're not dealing with cirq's type issues
+
+
+@dataclass
+class SnRuntime(UnsafeOperatorRuntimeABC):
+    def num_qubits(self) -> int:
+        return 1
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        if adjoint:
+            return SpRuntime().unsafe_apply(qubits, adjoint=False)
+
+        return [(cirq.X(*qubits) + 1j * cirq.Y(*qubits)) / 2]  # type: ignore  -- we're not dealing with cirq's type issues
+
+
+@dataclass
+class MultRuntime(OperatorRuntimeABC):
+    lhs: OperatorRuntimeABC
+    rhs: OperatorRuntimeABC
+
+    def num_qubits(self) -> int:
+        n = self.lhs.num_qubits()
+        assert n == self.rhs.num_qubits()
+        return n
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        rhs = self.rhs.unsafe_apply(qubits, adjoint=adjoint)
+        lhs = self.lhs.unsafe_apply(qubits, adjoint=adjoint)
+
+        if adjoint:
+            return lhs + rhs
+        else:
+            return rhs + lhs
+
+
+@dataclass
+class KronRuntime(OperatorRuntimeABC):
+    lhs: OperatorRuntimeABC
+    rhs: OperatorRuntimeABC
+
+    def num_qubits(self) -> int:
+        return self.lhs.num_qubits() + self.rhs.num_qubits()
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        n = self.lhs.num_qubits()
+        cirq_ops = self.lhs.unsafe_apply(qubits[:n], adjoint=adjoint)
+        cirq_ops.extend(self.rhs.unsafe_apply(qubits[n:], adjoint=adjoint))
+        return cirq_ops
+
+
+@dataclass
+class ControlRuntime(OperatorRuntimeABC):
+    operator: OperatorRuntimeABC
+    n_controls: int
+
+    def num_qubits(self) -> int:
+        return self.n_controls + self.operator.num_qubits()
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        m = len(qubits) - self.n_controls
+        cirq_ops = self.operator.unsafe_apply(qubits[m:], adjoint=adjoint)
+        controlled_ops = [cirq_op.controlled_by(*qubits[:m]) for cirq_op in cirq_ops]
+        return controlled_ops
+
+
+@dataclass
+class AdjointRuntime(OperatorRuntimeABC):
+    operator: OperatorRuntimeABC
+
+    def num_qubits(self) -> int:
+        return self.operator.num_qubits()
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        # NOTE: to account for e.g. adjoint(adjoint(op))
+        passed_on_adjoint = not adjoint
+        return self.operator.unsafe_apply(qubits, adjoint=passed_on_adjoint)
+
+
+@dataclass
+class U3Runtime(UnsafeOperatorRuntimeABC):
+    theta: float
+    phi: float
+    lam: float
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def angles(self, adjoint: bool) -> tuple[float, float, float]:
+        if adjoint:
+            # NOTE: adjoint(U(theta, phi, lam)) == U(-theta, -lam, -phi)
+            return -self.theta, -self.lam, -self.phi
+        else:
+            return self.theta, self.phi, self.lam
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        theta, phi, lam = self.angles(adjoint=adjoint)
+
+        ops = [
+            cirq.Rz(rads=lam)(*qubits),
+            cirq.Rx(rads=math.pi / 2)(*qubits),
+            cirq.Rz(rads=theta)(*qubits),
+            cirq.Rx(rads=-math.pi / 2)(*qubits),
+            cirq.Rz(rads=phi)(*qubits),
+        ]
+
+        return ops
+
+
+@dataclass
+class ScaleRuntime(OperatorRuntimeABC):
+    factor: Number
+    operator: OperatorRuntimeABC
+
+    def num_qubits(self) -> int:
+        return self.operator.num_qubits()
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        cirq_ops = self.operator.unsafe_apply(qubits=qubits, adjoint=adjoint)
+        return [self.factor * cirq_ops[0]] + cirq_ops[1:]  # type: ignore
+
+
+@dataclass
+class PauliStringRuntime(OperatorRuntimeABC):
+    string: str
+
+    def num_qubits(self) -> int:
+        return len(self.string)
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        pauli_mapping = {
+            qbit: pauli_label for (qbit, pauli_label) in zip(qubits, self.string)
+        }
+        return [cirq.PauliString(pauli_mapping)]

bloqade/squin/cirq/lowering.py

@@ -5,7 +5,7 @@ from dataclasses import field, dataclass
 import cirq
 from kirin import ir, lowering
 from kirin.rewrite import Walk, CFGCompactify
-from kirin.dialects import py, ilist
+from kirin.dialects import py, scf, ilist
 
 from .. import op, noise, qubit
 
@@ -150,10 +150,79 @@ class Squin(lowering.LoweringABC[CirqNode]):
     ):
         if len(node.qubits) == 1:
             qbit = self.lower_qubit_getindex(state, node.qubits[0])
-            return state.current_frame.push(qubit.MeasureQubit(qbit))
+            stmt = state.current_frame.push(qubit.MeasureQubit(qbit))
+        else:
+            qbits = self.lower_qubit_getindices(state, node.qubits)
+            stmt = state.current_frame.push(qubit.MeasureQubitList(qbits))
 
-        qbits = self.lower_qubit_getindices(state, node.qubits)
-        return state.current_frame.push(qubit.MeasureQubitList(qbits))
+        key = node.gate.key
+        if isinstance(key, cirq.MeasurementKey):
+            key = key.name
+
+        state.current_frame.defs[key] = stmt.result
+        return stmt
+
+    def visit_ClassicallyControlledOperation(
+        self, state: lowering.State[CirqNode], node: cirq.ClassicallyControlledOperation
+    ):
+        conditions: list[ir.SSAValue] = []
+        for outcome in node.classical_controls:
+            key = outcome.key
+            if isinstance(key, cirq.MeasurementKey):
+                key = key.name
+            measurement_outcome = state.current_frame.defs[key]
+
+            if measurement_outcome.type.is_subseteq(ilist.IListType):
+                # NOTE: there is currently no convenient ilist.any method, so we need to use foldl
+                # with a simple function that just does an or
+
+                def bool_op_or(x: bool, y: bool) -> bool:
+                    return x or y
+
+                f_code = state.current_frame.push(
+                    lowering.Python(self.dialects).python_function(bool_op_or)
+                )
+                fn = ir.Method(
+                    mod=None,
+                    py_func=bool_op_or,
+                    sym_name="bool_op_or",
+                    arg_names=[],
+                    dialects=self.dialects,
+                    code=f_code,
+                )
+                f_const = state.current_frame.push(py.constant.Constant(fn))
+                init_val = state.current_frame.push(py.Constant(False)).result
+                condition = state.current_frame.push(
+                    ilist.Foldl(f_const.result, measurement_outcome, init=init_val)
+                ).result
+            else:
+                condition = measurement_outcome
+
+            conditions.append(condition)
+
+        if len(conditions) == 1:
+            condition = conditions[0]
+        else:
+            condition = state.current_frame.push(
+                py.boolop.And(conditions[0], conditions[1])
+            ).result
+            for next_cond in conditions[2:]:
+                condition = state.current_frame.push(
+                    py.boolop.And(condition, next_cond)
+                ).result
+
+        then_stmt = self.visit(state, node.without_classical_controls())
+
+        assert isinstance(
+            then_stmt, ir.Statement
+        ), f"Expected operation of classically controlled node {node} to be lowered to a statement, got type {type(then_stmt)}. \
+        Please report this issue!"
+
+        # NOTE: remove stmt from parent block
+        then_stmt.detach()
+        then_body = ir.Block((then_stmt,))
+
+        return state.current_frame.push(scf.IfElse(condition, then_body=then_body))
 
     def visit_SingleQubitPauliStringGateOperation(
         self,

{bloqade_circuit-0.4.1.dist-info → bloqade_circuit-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.4.1
+Version: 0.4.3
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE

{bloqade_circuit-0.4.1.dist-info → bloqade_circuit-0.4.3.dist-info}/RECORD

@@ -74,7 +74,7 @@ bloqade/qasm2/emit/impls/noise.py,sha256=-N9PmCbz8MwC6xtd55GOpjDoWMyJPJBMVDWT3G8
 bloqade/qasm2/parse/__init__.py,sha256=01tlLfrR015nAAPWw3i_Cs9IXsShpXMnJMagcJ_Vuik,986
 bloqade/qasm2/parse/ast.py,sha256=a48ssf0D_xaE-27PsyeBD5lBvwN2Dojj-RWIBhy7jJE,2904
 bloqade/qasm2/parse/build.py,sha256=2CibD1ZRX3_aknmhb5XvFQcI2sBOn97DlQHomb9CMEw,10621
-bloqade/qasm2/parse/lowering.py,sha256=CwwCdaatbc5CiSmVcltGko2kblX60COuJDzqQtv-L4E,21274
+bloqade/qasm2/parse/lowering.py,sha256=VRFhdDsgAeALVzJxmDtlRI3L1xLrAGnz_Zh-qNOChhM,20844
 bloqade/qasm2/parse/parser.py,sha256=fxqp65dv8NnXE-Ie7ryLESfSH3Xr0unx1EBQysctiHM,121
 bloqade/qasm2/parse/print.py,sha256=PaigQ5RbcfhOteWvDdQHoKsTE3tcNefpVfh1sp5eZEI,8973
 bloqade/qasm2/parse/qasm2.lark,sha256=IYrBydUoVLn1VCNDPP5uNN5BHDET3fQ2yG11cOy900k,2238
@@ -117,8 +117,12 @@ bloqade/squin/analysis/nsites/__init__.py,sha256=RlQg7ivczXCXG5lMeL3ipYKj2oJKC4T
 bloqade/squin/analysis/nsites/analysis.py,sha256=rIe1RU1MZRItcE2aB8DYahLrv73HfD3IHCX3E_EGQ1c,1773
 bloqade/squin/analysis/nsites/impls.py,sha256=OaKuAoZ0EAorStYDZxzgc6Dk42kuj19MLkqHWG1MEQM,2592
 bloqade/squin/analysis/nsites/lattice.py,sha256=ruh0808SHtj3ecuT-C3AZTsLY2j3DRhtezGiTZvcuVs,942
-bloqade/squin/cirq/__init__.py,sha256=aGRwIWqbbKaW-1Pm697uo-wQsM9ZBf0h5Qlo3Y6HJ-Y,2453
-bloqade/squin/cirq/lowering.py,sha256=Crq4oq0O_ucZCjSXV6DNEMvZCo5HCl9Fy0mpYfYOdok,11689
+bloqade/squin/cirq/__init__.py,sha256=AptJlelH-KJoFKLnq6phq68SrV785zWzi2NOfLH62ms,5994
+bloqade/squin/cirq/lowering.py,sha256=4-kZFH_qbBbV-c3-C9KhIB5o_cp_D8oxJrS8KicD_A8,14382
+bloqade/squin/cirq/emit/emit_circuit.py,sha256=7puJ3eCFwE9VdPb9NAiSdyRNkoQPwo_uVykz9Yv7c14,3761
+bloqade/squin/cirq/emit/op.py,sha256=z54NP5KqMxffXeFGWamEzvunpTNrxmYuluurk4j2-ps,4000
+bloqade/squin/cirq/emit/qubit.py,sha256=Z2HUsZmJ5F2uHCPGru81ux2usoX77KwtS97_cgeJRMI,1910
+bloqade/squin/cirq/emit/runtime.py,sha256=6_oHod-WK5yv0ae9xziQn-eh4Hn3MZNNqu4kJtOzPeY,6543
 bloqade/squin/noise/__init__.py,sha256=HQl3FE0SZAGEX3qdveapCaMX391lgLvWeWnoE6Z2pYw,332
 bloqade/squin/noise/_dialect.py,sha256=2IR98J-lXm5Y3srP9g-FD4JC-qTq2seureM6mKKq1xg,63
 bloqade/squin/noise/_wrapper.py,sha256=0jD5va_go9jEW5rC6bZSWU30kjCha2-axFogPON3-V0,580
@@ -191,7 +195,7 @@ bloqade/visual/animation/runtime/atoms.py,sha256=EmjxhujLiHHPS_HtH_B-7TiqeHgvW5u
 bloqade/visual/animation/runtime/ppoly.py,sha256=JB9IP53N1w6adBJEue6J5Nmj818Id9JvrlgrmiQTU1I,1385
 bloqade/visual/animation/runtime/qpustate.py,sha256=rlmxQeJSvaohXrTpXQL5y-NJcpvfW33xPaYM1slv7cc,4270
 bloqade/visual/animation/runtime/utils.py,sha256=ju9IzOWX-vKwfpqUjlUKu3Ssr_UFPFFq-tzH_Nqyo_c,1212
-bloqade_circuit-0.4.1.dist-info/METADATA,sha256=53teyQGCbiQ0-uxPqwBrNAKdeLcpe7h3Aj7fgVq26HU,3683
-bloqade_circuit-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bloqade_circuit-0.4.1.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
-bloqade_circuit-0.4.1.dist-info/RECORD,,
+bloqade_circuit-0.4.3.dist-info/METADATA,sha256=ZBwC4D0_sLjDceh6og0S4tzzFghL5gTxP76QQrMsvv4,3683
+bloqade_circuit-0.4.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bloqade_circuit-0.4.3.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
+bloqade_circuit-0.4.3.dist-info/RECORD,,