bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/stim/dialects/aux/emit.py

@@ -0,0 +1,102 @@
+from kirin.emit import EmitStrFrame
+from kirin.interp import MethodTable, impl
+
+from bloqade.stim.emit.stim import EmitStimMain
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="emit.stim")
+class EmitStimAuxMethods(MethodTable):
+
+    @impl(stmts.ConstInt)
+    def const_int(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstInt):
+
+        out: str = f"{stmt.value}"
+
+        return (out,)
+
+    @impl(stmts.ConstFloat)
+    def const_float(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstFloat
+    ):
+
+        out: str = f"{stmt.value:.8f}"
+
+        return (out,)
+
+    @impl(stmts.ConstBool)
+    def const_bool(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstBool
+    ):
+        out: str = "!" if stmt.value else ""
+
+        return (out,)
+
+    @impl(stmts.ConstStr)
+    def const_str(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstBool):
+
+        return (stmt.value,)
+
+    @impl(stmts.Neg)
+    def neg(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Neg):
+
+        operand: str = frame.get(stmt.operand)
+
+        return ("-" + operand,)
+
+    @impl(stmts.GetRecord)
+    def get_rec(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.GetRecord):
+
+        id: str = frame.get(stmt.id)
+        out: str = f"rec[{id}]"
+
+        return (out,)
+
+    @impl(stmts.Tick)
+    def tick(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Tick):
+
+        emit.writeln(frame, "TICK")
+
+        return ()
+
+    @impl(stmts.Detector)
+    def detector(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Detector):
+
+        coords: tuple[str, ...] = frame.get_values(stmt.coord)
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+
+        coord_str: str = ", ".join(coords)
+        target_str: str = " ".join(targets)
+        emit.writeln(frame, f"DETECTOR({coord_str}) {target_str}")
+
+        return ()
+
+    @impl(stmts.ObservableInclude)
+    def obs_include(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ObservableInclude
+    ):
+
+        idx: str = frame.get(stmt.idx)
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+
+        target_str: str = " ".join(targets)
+        emit.writeln(frame, f"OBSERVABLE_INCLUDE({idx}) {target_str}")
+
+        return ()
+
+    @impl(stmts.NewPauliString)
+    def new_paulistr(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.NewPauliString
+    ):
+
+        string: tuple[str, ...] = frame.get_values(stmt.string)
+        flipped: tuple[str, ...] = frame.get_values(stmt.flipped)
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+
+        out = "*".join(
+            f"{flp}{base}{tgt}" for flp, base, tgt in zip(flipped, string, targets)
+        )
+
+        return (out,)

bloqade/stim/dialects/aux/interp.py

@@ -0,0 +1,39 @@
+from kirin import interp
+
+from . import stmts
+from .types import RecordResult
+from ._dialect import dialect
+
+
+@dialect.register
+class StimAuxMethods(interp.MethodTable):
+
+    @interp.impl(stmts.ConstFloat)
+    @interp.impl(stmts.ConstInt)
+    @interp.impl(stmts.ConstBool)
+    @interp.impl(stmts.ConstStr)
+    def const(
+        self,
+        interpreter: interp.Interpreter,
+        frame: interp.Frame,
+        stmt: stmts.ConstFloat | stmts.ConstInt | stmts.ConstBool | stmts.ConstStr,
+    ):
+        return (stmt.value,)
+
+    @interp.impl(stmts.Neg)
+    def neg(
+        self,
+        interpreter: interp.Interpreter,
+        frame: interp.Frame,
+        stmt: stmts.Neg,
+    ):
+        return (-frame.get(stmt.operand),)
+
+    @interp.impl(stmts.GetRecord)
+    def get_rec(
+        self,
+        interpreter: interp.Interpreter,
+        frame: interp.Frame,
+        stmt: stmts.GetRecord,
+    ):
+        return (RecordResult(value=frame.get(stmt.id)),)

bloqade/stim/dialects/aux/lowering.py

@@ -0,0 +1,40 @@
+import ast
+
+from kirin import lowering
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class StimAuxLowering(lowering.FromPythonAST):
+
+    def _const_stmt(
+        self, state: lowering.State, value: int | float | str | bool
+    ) -> stmts.ConstInt | stmts.ConstFloat | stmts.ConstStr | stmts.ConstBool:
+
+        if isinstance(value, bool):
+            return stmts.ConstBool(value=value)
+        elif isinstance(value, int):
+            return stmts.ConstInt(value=value)
+        elif isinstance(value, float):
+            return stmts.ConstFloat(value=value)
+        elif isinstance(value, str):
+            return stmts.ConstStr(value=value)
+        else:
+            raise lowering.BuildError(f"unsupported Stim constant type {type(value)}")
+
+    def lower_Constant(self, state: lowering.State, node: ast.Constant):
+        stmt = self._const_stmt(state, node.value)
+        return state.current_frame.push(stmt)
+
+    def lower_Expr(self, state: lowering.State, node: ast.Expr):
+        return state.parent.visit(state, node.value)  # just forward the visit
+
+    def lower_UnaryOp(self, state: lowering.State, node: ast.UnaryOp):
+        if isinstance(node.op, ast.USub):
+            value = state.lower(node.operand).expect_one()
+            stmt = stmts.Neg(operand=value)
+            return state.current_frame.push(stmt)
+        else:
+            raise lowering.BuildError(f"unsupported Stim unaryop {node.op}")

bloqade/stim/dialects/aux/stmts/__init__.py

@@ -0,0 +1,14 @@
+from .const import (
+    Neg as Neg,
+    ConstInt as ConstInt,
+    ConstStr as ConstStr,
+    ConstBool as ConstBool,
+    ConstFloat as ConstFloat,
+)
+from .annotate import (
+    Tick as Tick,
+    Detector as Detector,
+    GetRecord as GetRecord,
+    NewPauliString as NewPauliString,
+    ObservableInclude as ObservableInclude,
+)

bloqade/stim/dialects/aux/stmts/annotate.py

@@ -0,0 +1,47 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from ..types import RecordType, PauliStringType
+from .._dialect import dialect
+
+PyNum = types.Union(types.Int, types.Float)
+
+
+@statement(dialect=dialect)
+class GetRecord(ir.Statement):
+    name = "get_rec"
+    traits = frozenset({lowering.FromPythonCall()})
+    id: ir.SSAValue = info.argument(type=types.Int)
+    result: ir.ResultValue = info.result(type=RecordType)
+
+
+@statement(dialect=dialect)
+class Detector(ir.Statement):
+    name = "detector"
+    traits = frozenset({lowering.FromPythonCall()})
+    coord: tuple[ir.SSAValue, ...] = info.argument(PyNum)
+    targets: tuple[ir.SSAValue, ...] = info.argument(RecordType)
+
+
+@statement(dialect=dialect)
+class ObservableInclude(ir.Statement):
+    name = "obs.include"
+    traits = frozenset({lowering.FromPythonCall()})
+    idx: ir.SSAValue = info.argument(type=types.Int)
+    targets: tuple[ir.SSAValue, ...] = info.argument(RecordType)
+
+
+@statement(dialect=dialect)
+class Tick(ir.Statement):
+    name = "tick"
+    traits = frozenset({lowering.FromPythonCall()})
+
+
+@statement(dialect=dialect)
+class NewPauliString(ir.Statement):
+    name = "new_pauli_string"
+    traits = frozenset({lowering.FromPythonCall()})
+    string: tuple[ir.SSAValue, ...] = info.argument(types.String)
+    flipped: tuple[ir.SSAValue, ...] = info.argument(types.Bool)
+    targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
+    result: ir.ResultValue = info.result(type=PauliStringType)

bloqade/stim/dialects/aux/stmts/const.py

@@ -0,0 +1,95 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.print import Printer
+
+from .._dialect import dialect as dialect
+
+
+@statement(dialect=dialect)
+class ConstInt(ir.Statement):
+    """IR Statement representing a constant integer value."""
+
+    name = "constant.int"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})
+    value: int = info.attribute(types.Int)
+    """value (int): The constant integer value."""
+    result: ir.ResultValue = info.result(types.Int)
+    """result (Int): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class ConstFloat(ir.Statement):
+    """IR Statement representing a constant float value."""
+
+    name = "constant.float"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})
+    value: float = info.attribute(types.Float)
+    """value (float): The constant float value."""
+    result: ir.ResultValue = info.result(types.Float)
+    """result (Float): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class ConstBool(ir.Statement):
+    """IR Statement representing a constant float value."""
+
+    name = "constant.bool"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})
+    value: bool = info.attribute(types.Bool)
+    """value (float): The constant float value."""
+    result: ir.ResultValue = info.result(types.Bool)
+    """result (Float): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class ConstStr(ir.Statement):
+    """IR Statement representing a constant str value."""
+
+    name = "constant.str"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})
+    value: str = info.attribute(types.String)
+    """value (str): The constant str value."""
+    result: ir.ResultValue = info.result(types.String)
+    """result (str): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class Neg(ir.Statement):
+    """IR Statement representing a negation operation."""
+
+    name = "neg"
+    traits = frozenset({lowering.FromPythonCall()})
+    operand: ir.SSAValue = info.argument(types.Int)
+    result: ir.ResultValue = info.result(types.Int)

bloqade/stim/dialects/aux/types.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from kirin import types
+
+
+@dataclass
+class RecordResult:
+    value: int
+
+
+@dataclass
+class PauliString:
+    string: tuple[str, ...]
+    flipped: tuple[bool, ...]
+    targets: tuple[int, ...]
+
+
+RecordType = types.PyClass(RecordResult)
+PauliStringType = types.PyClass(PauliString)

bloqade/stim/dialects/collapse/__init__.py

@@ -0,0 +1,3 @@
+from .emit import EmitStimCollapseMethods as EmitStimCollapseMethods
+from .stmts import *  # noqa F403
+from ._dialect import dialect as dialect

bloqade/stim/dialects/collapse/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("stim.collapse")

bloqade/stim/dialects/collapse/emit.py

@@ -0,0 +1,68 @@
+from kirin.emit import EmitStrFrame
+from kirin.interp import MethodTable, impl
+
+from bloqade.stim.emit.stim import EmitStimMain
+
+from . import stmts
+from ._dialect import dialect
+from .stmts.reset import Reset
+from .stmts.measure import Measurement
+
+
+@dialect.register(key="emit.stim")
+class EmitStimCollapseMethods(MethodTable):
+
+    meas_map: dict[str, str] = {
+        stmts.MX.name: "MX",
+        stmts.MY.name: "MY",
+        stmts.MZ.name: "MZ",
+        stmts.MXX.name: "MXX",
+        stmts.MYY.name: "MYY",
+        stmts.MZZ.name: "MZZ",
+    }
+
+    @impl(stmts.MX)
+    @impl(stmts.MY)
+    @impl(stmts.MZ)
+    @impl(stmts.MXX)
+    @impl(stmts.MYY)
+    @impl(stmts.MZZ)
+    def get_measure(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: Measurement):
+
+        probability: str = frame.get(stmt.p)
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+
+        out = f"{self.meas_map[stmt.name]}({probability}) " + " ".join(targets)
+        emit.writeln(frame, out)
+
+        return ()
+
+    reset_map: dict[str, str] = {
+        stmts.RX.name: "RX",
+        stmts.RY.name: "RY",
+        stmts.RZ.name: "RZ",
+    }
+
+    @impl(stmts.RX)
+    @impl(stmts.RY)
+    @impl(stmts.RZ)
+    def get_reset(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: Reset):
+
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+
+        out = f"{self.reset_map[stmt.name]} " + " ".join(targets)
+        emit.writeln(frame, out)
+
+        return ()
+
+    @impl(stmts.PPMeasurement)
+    def pp_measure(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.PPMeasurement
+    ):
+        probability: str = frame.get(stmt.p)
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+
+        out = f"MPP({probability}) " + " ".join(targets)
+        emit.writeln(frame, out)
+
+        return ()

bloqade/stim/dialects/collapse/stmts/__init__.py

@@ -0,0 +1,3 @@
+from .reset import RX as RX, RY as RY, RZ as RZ
+from .measure import MX as MX, MY as MY, MZ as MZ, MXX as MXX, MYY as MYY, MZZ as MZZ
+from .pp_measure import PPMeasurement as PPMeasurement

bloqade/stim/dialects/collapse/stmts/measure.py

@@ -0,0 +1,45 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from .._dialect import dialect
+
+
+@statement
+class Measurement(ir.Statement):
+    name = "measurement"
+    traits = frozenset({lowering.FromPythonCall()})
+    p: ir.SSAValue = info.argument(types.Float)
+    """probability of noise introduced by measurement. For example 0.01 means 1% the measurement will be flipped"""
+    targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
+
+
+# 1Q measurements
+@statement(dialect=dialect)
+class MZ(Measurement):
+    name = "MZ"
+
+
+@statement(dialect=dialect)
+class MY(Measurement):
+    name = "MY"
+
+
+@statement(dialect=dialect)
+class MX(Measurement):
+    name = "MX"
+
+
+# Pair measurements
+@statement(dialect=dialect)
+class MZZ(Measurement):
+    name = "MZZ"
+
+
+@statement(dialect=dialect)
+class MYY(Measurement):
+    name = "MYY"
+
+
+@statement(dialect=dialect)
+class MXX(Measurement):
+    name = "MXX"

bloqade/stim/dialects/collapse/stmts/pp_measure.py

@@ -0,0 +1,14 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from .._dialect import dialect
+from ...aux.types import PauliStringType
+
+
+@statement(dialect=dialect)
+class PPMeasurement(ir.Statement):
+    name = "MPP"
+    traits = frozenset({lowering.FromPythonCall()})
+    p: ir.SSAValue = info.argument(types.Float)
+    """probability of noise introduced by measurement. For example 0.01 means 1% the measurement will be flipped"""
+    targets: tuple[ir.SSAValue, ...] = info.argument(PauliStringType)

bloqade/stim/dialects/collapse/stmts/reset.py

@@ -0,0 +1,26 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from .._dialect import dialect
+
+
+@statement
+class Reset(ir.Statement):
+    name = "reset"
+    traits = frozenset({lowering.FromPythonCall()})
+    targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
+
+
+@statement(dialect=dialect)
+class RZ(Reset):
+    name = "RZ"
+
+
+@statement(dialect=dialect)
+class RY(Reset):
+    name = "RY"
+
+
+@statement(dialect=dialect)
+class RX(Reset):
+    name = "RX"

bloqade/stim/dialects/gate/__init__.py

@@ -0,0 +1,3 @@
+from .emit import EmitStimGateMethods as EmitStimGateMethods
+from .stmts import *  # noqa F403
+from ._dialect import dialect as dialect

bloqade/stim/dialects/gate/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("stim.gate")

bloqade/stim/dialects/gate/emit.py

@@ -0,0 +1,87 @@
+from kirin.emit import EmitStrFrame
+from kirin.interp import MethodTable, impl
+
+from bloqade.stim.emit.stim import EmitStimMain
+
+from . import stmts
+from ._dialect import dialect
+from .stmts.base import SingleQubitGate, ControlledTwoQubitGate
+
+
+@dialect.register(key="emit.stim")
+class EmitStimGateMethods(MethodTable):
+
+    gate_1q_map: dict[str, tuple[str, str]] = {
+        stmts.X.name: ("X", "X"),
+        stmts.Y.name: ("Y", "Y"),
+        stmts.Z.name: ("Z", "Z"),
+        stmts.H.name: ("H", "H"),
+        stmts.S.name: ("S", "S_DAG"),
+        stmts.SqrtX.name: ("SQRT_X", "SQRT_X_DAG"),
+        stmts.SqrtY.name: ("SQRT_Y", "SQRT_Y_DAG"),
+        stmts.SqrtZ.name: ("SQRT_Z", "SQRT_Z_DAG"),
+    }
+
+    @impl(stmts.X)
+    @impl(stmts.Y)
+    @impl(stmts.Z)
+    @impl(stmts.S)
+    @impl(stmts.H)
+    @impl(stmts.SqrtX)
+    @impl(stmts.SqrtY)
+    @impl(stmts.SqrtZ)
+    def single_qubit_gate(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: SingleQubitGate
+    ):
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        res = f"{self.gate_1q_map[stmt.name][int(stmt.dagger)]} " + " ".join(targets)
+        emit.writeln(frame, res)
+
+        return ()
+
+    gate_2q_map: dict[str, tuple[str, str]] = {
+        stmts.Swap.name: ("SWAP", "SWAP"),
+    }
+
+    @impl(stmts.Swap)
+    def two_qubit_gate(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: ControlledTwoQubitGate
+    ):
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        res = f"{self.gate_ctrl_2q_map[stmt.name][int(stmt.dagger)]} " + " ".join(
+            targets
+        )
+        emit.writeln(frame, res)
+
+        return ()
+
+    gate_ctrl_2q_map: dict[str, tuple[str, str]] = {
+        stmts.CX.name: ("CX", "CX"),
+        stmts.CY.name: ("CY", "CY"),
+        stmts.CZ.name: ("CZ", "CZ"),
+        stmts.Swap.name: ("SWAP", "SWAP"),
+    }
+
+    @impl(stmts.CX)
+    @impl(stmts.CY)
+    @impl(stmts.CZ)
+    def ctrl_two_qubit_gate(
+        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: ControlledTwoQubitGate
+    ):
+        controls: tuple[str, ...] = frame.get_values(stmt.controls)
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        res = f"{self.gate_ctrl_2q_map[stmt.name][int(stmt.dagger)]} " + " ".join(
+            f"{ctrl} {tgt}" for ctrl, tgt in zip(controls, targets)
+        )
+        emit.writeln(frame, res)
+
+        return ()
+
+    @impl(stmts.SPP)
+    def spp(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.SPP):
+
+        targets: tuple[str, ...] = frame.get_values(stmt.targets)
+        res = "SPP " + " ".join(targets)
+        emit.writeln(frame, res)
+
+        return ()

bloqade/stim/dialects/gate/stmts/__init__.py

@@ -0,0 +1,14 @@
+from .pp import SPP as SPP
+from .control_2q import CX as CX, CY as CY, CZ as CZ
+from .clifford_1q import (
+    H as H,
+    S as S,
+    X as X,
+    Y as Y,
+    Z as Z,
+    SqrtX as SqrtX,
+    SqrtY as SqrtY,
+    SqrtZ as SqrtZ,
+    Identity as Identity,
+)
+from .clifford_2q import Swap as Swap

bloqade/stim/dialects/gate/stmts/base.py

@@ -0,0 +1,31 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from bloqade.stim.dialects.aux import RecordType
+
+
+@statement
+class Gate(ir.Statement):
+    name = "stim_gate"
+    traits = frozenset({lowering.FromPythonCall()})
+    targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
+    dagger: bool = info.attribute(default=False)
+
+
+@statement
+class SingleQubitGate(Gate):
+    name = "single_qubit_gate"
+
+
+@statement
+class TwoQubitGate(Gate):
+    name = "two_qubit_gate"
+
+
+# control can either be a qubit or a measurement record
+@statement
+class ControlledTwoQubitGate(Gate):
+    name = "controlled_two_qubit_gate"
+    controls: tuple[ir.SSAValue, ...] = info.argument(
+        types.Union(types.Int, RecordType)
+    )