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

bloqade/squin/cirq/lowering.py

@@ -0,0 +1,303 @@
+import math
+from typing import Any
+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 .. import op, noise, qubit
+
+CirqNode = cirq.Circuit | cirq.Moment | cirq.Gate | cirq.Qid | cirq.Operation
+
+DecomposeNode = (
+    cirq.SwapPowGate
+    | cirq.ISwapPowGate
+    | cirq.PhasedXPowGate
+    | cirq.PhasedXZGate
+    | cirq.CSwapGate
+)
+
+
+@dataclass
+class Squin(lowering.LoweringABC[CirqNode]):
+    """Lower a cirq.Circuit object to a squin kernel"""
+
+    circuit: cirq.Circuit
+    qreg: qubit.New = field(init=False)
+    qreg_index: dict[cirq.Qid, int] = field(init=False, default_factory=dict)
+    next_qreg_index: int = field(init=False, default=0)
+
+    def lower_qubit_getindex(self, state: lowering.State[CirqNode], qid: cirq.Qid):
+        index = self.qreg_index.get(qid)
+
+        if index is None:
+            index = self.next_qreg_index
+            self.qreg_index[qid] = index
+            self.next_qreg_index += 1
+
+        index_ssa = state.current_frame.push(py.Constant(index)).result
+        qbit_getitem = state.current_frame.push(py.GetItem(self.qreg.result, index_ssa))
+        return qbit_getitem.result
+
+    def lower_qubit_getindices(
+        self, state: lowering.State[CirqNode], qids: list[cirq.Qid]
+    ):
+        qbits_getitem = [self.lower_qubit_getindex(state, qid) for qid in qids]
+        qbits_stmt = ilist.New(values=qbits_getitem)
+        qbits_result = state.current_frame.get(qbits_stmt.name)
+
+        if qbits_result is not None:
+            return qbits_result
+
+        state.current_frame.push(qbits_stmt)
+        return qbits_stmt.result
+
+    def run(
+        self,
+        stmt: CirqNode,
+        *,
+        source: str | None = None,
+        globals: dict[str, Any] | None = None,
+        file: str | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+    ) -> ir.Region:
+
+        state = lowering.State(
+            self,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+
+        with state.frame(
+            [stmt],
+            globals=globals,
+            finalize_next=False,
+        ) as frame:
+            # NOTE: create a global register of qubits first
+            # TODO: can there be a circuit without qubits?
+            n_qubits = cirq.num_qubits(self.circuit)
+            n = frame.push(py.Constant(n_qubits))
+            self.qreg = frame.push(qubit.New(n_qubits=n.result))
+
+            self.visit(state, stmt)
+
+            if compactify:
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
+
+            region = frame.curr_region
+
+        return region
+
+    def visit(self, state: lowering.State[CirqNode], node: CirqNode) -> lowering.Result:
+        name = node.__class__.__name__
+        return getattr(self, f"visit_{name}", self.generic_visit)(state, node)
+
+    def generic_visit(self, state: lowering.State[CirqNode], node: CirqNode):
+        if isinstance(node, CirqNode):
+            raise lowering.BuildError(
+                f"Cannot lower {node.__class__.__name__} node: {node}"
+            )
+        raise lowering.BuildError(
+            f"Unexpected `{node.__class__.__name__}` node: {repr(node)} is not an AST node"
+        )
+
+    def lower_literal(self, state: lowering.State[CirqNode], value) -> ir.SSAValue:
+        raise lowering.BuildError("Literals not supported in cirq circuit")
+
+    def lower_global(
+        self, state: lowering.State[CirqNode], node: CirqNode
+    ) -> lowering.LoweringABC.Result:
+        raise lowering.BuildError("Literals not supported in cirq circuit")
+
+    def visit_Circuit(
+        self, state: lowering.State[CirqNode], node: cirq.Circuit
+    ) -> lowering.Result:
+        for moment in node:
+            state.lower(moment)
+
+    def visit_Moment(
+        self, state: lowering.State[CirqNode], node: cirq.Moment
+    ) -> lowering.Result:
+        for op_ in node.operations:
+            state.lower(op_)
+
+    def visit_GateOperation(
+        self, state: lowering.State[CirqNode], node: cirq.GateOperation
+    ):
+        if isinstance(node.gate, cirq.MeasurementGate):
+            # NOTE: special dispatch here, since measurement is a gate + a qubit in cirq,
+            # but a single statement in squin
+            return self.lower_measurement(state, node)
+
+        if isinstance(node.gate, DecomposeNode):
+            # NOTE: easier to decompose these, but for that we need the qubits too,
+            # so we need to do this within this method
+            for subnode in cirq.decompose_once(node):
+                state.lower(subnode)
+            return
+
+        op_ = state.lower(node.gate).expect_one()
+        qbits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(qubit.Apply(operator=op_, qubits=qbits))
+
+    def lower_measurement(
+        self, state: lowering.State[CirqNode], node: cirq.GateOperation
+    ):
+        if len(node.qubits) == 1:
+            qbit = self.lower_qubit_getindex(state, node.qubits[0])
+            return state.current_frame.push(qubit.MeasureQubit(qbit))
+
+        qbits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(qubit.MeasureQubitList(qbits))
+
+    def visit_SingleQubitPauliStringGateOperation(
+        self,
+        state: lowering.State[CirqNode],
+        node: cirq.SingleQubitPauliStringGateOperation,
+    ):
+
+        match node.pauli:
+            case cirq.X:
+                op_ = op.stmts.X()
+            case cirq.Y:
+                op_ = op.stmts.Y()
+            case cirq.Z:
+                op_ = op.stmts.Z()
+            case cirq.I:
+                op_ = op.stmts.Identity(sites=1)
+            case _:
+                raise lowering.BuildError(f"Unexpected Pauli operation {node.pauli}")
+
+        state.current_frame.push(op_)
+        qargs = self.lower_qubit_getindices(state, [node.qubit])
+        return state.current_frame.push(qubit.Apply(op_.result, qargs))
+
+    def visit_HPowGate(self, state: lowering.State[CirqNode], node: cirq.HPowGate):
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.H())
+
+        return state.lower(node.in_su2())
+
+    def visit_XPowGate(self, state: lowering.State[CirqNode], node: cirq.XPowGate):
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.X())
+
+        return self.visit(state, node.in_su2())
+
+    def visit_YPowGate(self, state: lowering.State[CirqNode], node: cirq.YPowGate):
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.Y())
+
+        return self.visit(state, node.in_su2())
+
+    def visit_ZPowGate(self, state: lowering.State[CirqNode], node: cirq.ZPowGate):
+        if node.exponent == 0.5:
+            return state.current_frame.push(op.stmts.S())
+
+        if node.exponent == 0.25:
+            return state.current_frame.push(op.stmts.T())
+
+        if node.exponent == 1:
+            return state.current_frame.push(op.stmts.Z())
+
+        # NOTE: just for the Z gate, an arbitrary exponent is equivalent to the ShiftOp
+        t = node.exponent
+        theta = state.current_frame.push(py.Constant(math.pi * t))
+        return state.current_frame.push(op.stmts.ShiftOp(theta=theta.result))
+
+    def visit_Rx(self, state: lowering.State[CirqNode], node: cirq.Rx):
+        x = state.current_frame.push(op.stmts.X())
+        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
+        return state.current_frame.push(op.stmts.Rot(axis=x.result, angle=angle.result))
+
+    def visit_Ry(self, state: lowering.State[CirqNode], node: cirq.Ry):
+        y = state.current_frame.push(op.stmts.Y())
+        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
+        return state.current_frame.push(op.stmts.Rot(axis=y.result, angle=angle.result))
+
+    def visit_Rz(self, state: lowering.State[CirqNode], node: cirq.Rz):
+        z = state.current_frame.push(op.stmts.Z())
+        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
+        return state.current_frame.push(op.stmts.Rot(axis=z.result, angle=angle.result))
+
+    def visit_CXPowGate(self, state: lowering.State[CirqNode], node: cirq.CXPowGate):
+        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(x, n_controls=1))
+
+    def visit_CZPowGate(self, state: lowering.State[CirqNode], node: cirq.CZPowGate):
+        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(z, n_controls=1))
+
+    def visit_ControlledOperation(
+        self, state: lowering.State[CirqNode], node: cirq.ControlledOperation
+    ):
+        return self.visit_GateOperation(state, node)
+
+    def visit_ControlledGate(
+        self, state: lowering.State[CirqNode], node: cirq.ControlledGate
+    ):
+        op_ = state.lower(node.sub_gate).expect_one()
+        n_controls = node.num_controls()
+        return state.current_frame.push(op.stmts.Control(op_, n_controls=n_controls))
+
+    def visit_XXPowGate(self, state: lowering.State[CirqNode], node: cirq.XXPowGate):
+        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Kron(x, x))
+
+    def visit_YYPowGate(self, state: lowering.State[CirqNode], node: cirq.YYPowGate):
+        y = state.lower(cirq.YPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Kron(y, y))
+
+    def visit_ZZPowGate(self, state: lowering.State[CirqNode], node: cirq.ZZPowGate):
+        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Kron(z, z))
+
+    def visit_CCXPowGate(self, state: lowering.State[CirqNode], node: cirq.CCXPowGate):
+        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(x, n_controls=2))
+
+    def visit_CCZPowGate(self, state: lowering.State[CirqNode], node: cirq.CCZPowGate):
+        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
+        return state.current_frame.push(op.stmts.Control(z, n_controls=2))
+
+    def visit_BitFlipChannel(
+        self, state: lowering.State[CirqNode], node: cirq.BitFlipChannel
+    ):
+        x = state.current_frame.push(op.stmts.X())
+        p = state.current_frame.push(py.Constant(node.p))
+        return state.current_frame.push(
+            noise.stmts.PauliError(basis=x.result, p=p.result)
+        )
+
+    def visit_AmplitudeDampingChannel(
+        self, state: lowering.State[CirqNode], node: cirq.AmplitudeDampingChannel
+    ):
+        r = state.current_frame.push(op.stmts.Reset())
+        p = state.current_frame.push(py.Constant(node.gamma))
+
+        # TODO: do we need a dedicated noise stmt for this? Using PauliError
+        # with this basis feels like a hack
+        noise_channel = state.current_frame.push(
+            noise.stmts.PauliError(basis=r.result, p=p.result)
+        )
+
+        return noise_channel
+
+    def visit_GeneralizedAmplitudeDampingChannel(
+        self,
+        state: lowering.State[CirqNode],
+        node: cirq.GeneralizedAmplitudeDampingChannel,
+    ):
+        raise NotImplementedError("TODO: needs a new operator statement")
+        # p = state.current_frame.push(py.Constant(node.p))
+        # gamma = state.current_frame.push(py.Constant(node.gamma))
+
+        # p1 =
+
+        # x = state.current_frame.push(op.stmts.X())
+        # noise_channel1 = noise.stmts.PauliError(basis=x.result, p=)

bloqade/squin/groups.py

@@ -1,19 +1,19 @@
 from kirin import ir, passes
 from kirin.prelude import structural_no_opt
+from kirin.rewrite import Walk, Chain
 from kirin.dialects import ilist
-from kirin.rewrite.walk import Walk
 
-from . import op, wire, qubit
+from . import op, wire, noise, qubit
 from .op.rewrite import PyMultToSquinMult
-from .rewrite.measure_desugar import MeasureDesugarRule
+from .rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
 
 
-@ir.dialect_group(structural_no_opt.union([op, qubit]))
+@ir.dialect_group(structural_no_opt.union([op, qubit, noise]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
     ilist_desugar_pass = ilist.IListDesugar(self)
-    measure_desugar_pass = Walk(MeasureDesugarRule())
+    desugar_pass = Walk(Chain(MeasureDesugarRule(), ApplyDesugarRule()))
     py_mult_to_mult_pass = PyMultToSquinMult(self)
 
     def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
@@ -25,7 +25,7 @@ def kernel(self):
 
         if typeinfer:
             typeinfer_pass(method)
-            measure_desugar_pass.rewrite(method.code)
+            desugar_pass.rewrite(method.code)
 
         ilist_desugar_pass(method)
 
@@ -36,7 +36,7 @@ def kernel(self):
     return run_pass
 
 
-@ir.dialect_group(structural_no_opt.union([op, wire]))
+@ir.dialect_group(structural_no_opt.union([op, wire, noise]))
 def wired(self):
     py_mult_to_mult_pass = PyMultToSquinMult(self)
 

bloqade/squin/lowering.py

@@ -0,0 +1,27 @@
+import ast
+from dataclasses import dataclass
+
+from kirin import lowering
+
+from . import qubit
+
+
+@dataclass(frozen=True)
+class ApplyAnyCallLowering(lowering.FromPythonCall["qubit.ApplyAny"]):
+    """
+    Custom lowering for ApplyAny that collects vararg qubits into a single tuple argument
+    """
+
+    def lower(
+        self, stmt: type["qubit.ApplyAny"], state: lowering.State, node: ast.Call
+    ):
+        if len(node.args) < 2:
+            raise lowering.BuildError(
+                "Apply requires at least one operator and one qubit as arguments!"
+            )
+        op, *qubits = node.args
+        op_ssa = state.lower(op).expect_one()
+        qubits_lowered = [state.lower(qbit).expect_one() for qbit in qubits]
+
+        s = stmt(op_ssa, tuple(qubits_lowered))
+        return state.current_frame.push(s)

bloqade/squin/noise/__init__.py

@@ -4,5 +4,7 @@ from ._wrapper import (
     pp_error as pp_error,
     depolarize as depolarize,
     qubit_loss as qubit_loss,
-    pauli_channel as pauli_channel,
+    pauli_error as pauli_error,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
 )

bloqade/squin/noise/_wrapper.py

@@ -14,11 +14,15 @@ def pp_error(op: Op, p: float) -> Op: ...
 
 
 @wraps(stmts.Depolarize)
-def depolarize(n_qubits: int, p: float) -> Op: ...
+def depolarize(p: float) -> Op: ...
 
 
-@wraps(stmts.PauliChannel)
-def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
+@wraps(stmts.SingleQubitPauliChannel)
+def single_qubit_pauli_channel(params: tuple[float, float, float]) -> Op: ...
+
+
+@wraps(stmts.TwoQubitPauliChannel)
+def two_qubit_pauli_channel(params: tuple[float, ...]) -> Op: ...
 
 
 @wraps(stmts.QubitLoss)

bloqade/squin/noise/rewrite.py

@@ -0,0 +1,111 @@
+import itertools
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import Walk
+from kirin.dialects import ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from .stmts import (
+    PPError,
+    QubitLoss,
+    Depolarize,
+    PauliError,
+    NoiseChannel,
+    TwoQubitPauliChannel,
+    SingleQubitPauliChannel,
+    StochasticUnitaryChannel,
+)
+from ..op.stmts import X, Y, Z, Kron, Identity
+
+
+class _RewriteNoiseStmts(RewriteRule):
+    """Rewrites squin noise statements to StochasticUnitaryChannel"""
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, NoiseChannel) or isinstance(node, QubitLoss):
+            return RewriteResult()
+
+        return getattr(self, "rewrite_" + node.name)(node)
+
+    def rewrite_pauli_error(self, node: PauliError) -> RewriteResult:
+        (operators := ilist.New(values=(node.basis,))).insert_before(node)
+        (ps := ilist.New(values=(node.p,))).insert_before(node)
+        stochastic_channel = StochasticUnitaryChannel(
+            operators=operators.result, probabilities=ps.result
+        )
+
+        node.replace_by(stochastic_channel)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_single_qubit_pauli_channel(
+        self, node: SingleQubitPauliChannel
+    ) -> RewriteResult:
+        paulis = (X(), Y(), Z())
+        paulis_ssa: list[ir.SSAValue] = []
+        for op in paulis:
+            op.insert_before(node)
+            paulis_ssa.append(op.result)
+
+        (pauli_ops := ilist.New(values=paulis_ssa)).insert_before(node)
+
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=pauli_ops.result, probabilities=node.params
+        )
+        node.replace_by(stochastic_unitary)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_two_qubit_pauli_channel(
+        self, node: TwoQubitPauliChannel
+    ) -> RewriteResult:
+        paulis = (X(), Y(), Z(), Identity(sites=1))
+        for op in paulis:
+            op.insert_before(node)
+
+        # NOTE: collect list so we can skip the last entry, which will be two identities
+        combinations = list(itertools.product(paulis, repeat=2))[:-1]
+        operators: list[ir.SSAValue] = []
+        for pauli_1, pauli_2 in combinations:
+            op = Kron(pauli_1.result, pauli_2.result)
+            op.insert_before(node)
+            operators.append(op.result)
+
+        (operator_list := ilist.New(values=operators)).insert_before(node)
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list.result, probabilities=node.params
+        )
+
+        node.replace_by(stochastic_unitary)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_p_p_error(self, node: PPError) -> RewriteResult:
+        (operators := ilist.New(values=(node.op,))).insert_before(node)
+        (ps := ilist.New(values=(node.p,))).insert_before(node)
+        stochastic_channel = StochasticUnitaryChannel(
+            operators=operators.result, probabilities=ps.result
+        )
+
+        node.replace_by(stochastic_channel)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_depolarize(self, node: Depolarize) -> RewriteResult:
+        paulis = (X(), Y(), Z())
+        operators: list[ir.SSAValue] = []
+        for op in paulis:
+            op.insert_before(node)
+            operators.append(op.result)
+
+        (operator_list := ilist.New(values=operators)).insert_before(node)
+        (ps := ilist.New(values=[node.p for _ in range(3)])).insert_before(node)
+
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list.result, probabilities=ps.result
+        )
+        node.replace_by(stochastic_unitary)
+
+        return RewriteResult(has_done_something=True)
+
+
+class RewriteNoiseStmts(Pass):
+    def unsafe_run(self, mt: ir.Method):
+        return Walk(_RewriteNoiseStmts()).rewrite(mt.code)

bloqade/squin/noise/stmts.py

@@ -1,21 +1,23 @@
-from kirin import ir, types
+from kirin import ir, types, lowering
 from kirin.decl import info, statement
+from kirin.dialects import ilist
 
 from bloqade.squin.op.types import OpType
 
 from ._dialect import dialect
+from ..op.types import NumOperators
 
 
 @statement
 class NoiseChannel(ir.Statement):
-    pass
+    traits = frozenset({lowering.FromPythonCall()})
+    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
 class PauliError(NoiseChannel):
     basis: ir.SSAValue = info.argument(OpType)
     p: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
@@ -26,34 +28,37 @@ class PPError(NoiseChannel):
 
     op: ir.SSAValue = info.argument(OpType)
     p: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
 class Depolarize(NoiseChannel):
     """
-    Apply n-qubit depolaize error to qubits
-    NOTE For Stim, this can only accept 1 or 2 qubits
+    Apply depolarize error to qubit
     """
 
-    n_qubits: int = info.attribute(types.Int)
     p: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(OpType)
 
 
 @statement(dialect=dialect)
-class PauliChannel(NoiseChannel):
-    # NOTE:
-    # 1-qubit 3 params px, py, pz
-    # 2-qubit 15 params pix, piy, piz, pxi, pxx, pxy, pxz, pyi, pyx ..., pzz
-    # TODO add validation for params (maybe during lowering via custom lower?)
-    n_qubits: int = info.attribute()
-    params: ir.SSAValue = info.argument(types.Tuple[types.Vararg(types.Float)])
-    result: ir.ResultValue = info.result(OpType)
+class SingleQubitPauliChannel(NoiseChannel):
+    params: ir.SSAValue = info.argument(ilist.IListType[types.Float, types.Literal(3)])
+
+
+@statement(dialect=dialect)
+class TwoQubitPauliChannel(NoiseChannel):
+    params: ir.SSAValue = info.argument(ilist.IListType[types.Float, types.Literal(15)])
 
 
 @statement(dialect=dialect)
 class QubitLoss(NoiseChannel):
     # NOTE: qubit loss error (not supported by Stim)
     p: ir.SSAValue = info.argument(types.Float)
+
+
+@statement(dialect=dialect)
+class StochasticUnitaryChannel(ir.Statement):
+    operators: ir.SSAValue = info.argument(ilist.IListType[OpType, NumOperators])
+    probabilities: ir.SSAValue = info.argument(
+        ilist.IListType[types.Float, NumOperators]
+    )
     result: ir.ResultValue = info.result(OpType)

bloqade/squin/op/__init__.py

@@ -25,6 +25,7 @@ from ._wrapper import (
     kron as kron,
     mult as mult,
     phase as phase,
+    reset as reset,
     scale as scale,
     shift as shift,
     spin_n as spin_n,

bloqade/squin/op/_wrapper.py

@@ -37,6 +37,10 @@ def control(op: types.Op, *, n_controls: int) -> types.Op:
     ...
 
 
+@wraps(stmts.Reset)
+def reset() -> types.Op: ...
+
+
 @wraps(stmts.Identity)
 def identity(*, sites: int) -> types.Op: ...