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

bloqade/squin/stdlib/simple/noise.py

@@ -0,0 +1,126 @@
+from typing import Any, Literal, TypeVar
+
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit
+
+from .. import broadcast
+from ...groups import kernel
+
+
+@kernel
+def depolarize(p: float, qubit: Qubit) -> None:
+    """
+    Apply a depolarizing noise channel to a qubit with probability `p`.
+
+    This will randomly select one of the Pauli operators X, Y, Z
+    with a probability `p / 3` and apply it to the qubit. No operator is applied
+    with a probability of `1 - p`.
+
+    Args:
+        p (float): The probability with which a Pauli operator is applied.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    broadcast.depolarize(p, ilist.IList([qubit]))
+
+
+N = TypeVar("N", bound=int)
+
+
+@kernel
+def depolarize2(p: float, control: Qubit, target: Qubit) -> None:
+    """
+    Symmetric two-qubit depolarization channel applied to a pair of qubits.
+
+    This will randomly select one of the pauli products
+
+    `{IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}`
+
+    each with a probability `p / 15`. No noise is applied with a probability of `1 - p`.
+
+    Args:
+        p (float): The probability with which a Pauli product is applied.
+        control (Qubit): The control qubit.
+        target (Qubit): The target qubit.
+    """
+    broadcast.depolarize2(p, ilist.IList([control]), ilist.IList([target]))
+
+
+@kernel
+def single_qubit_pauli_channel(px: float, py: float, pz: float, qubit: Qubit) -> None:
+    """
+    Apply a Pauli error channel with weighted `px, py, pz`. No error is applied with a probability
+    `1 - (px + py + pz)`.
+
+    This randomly selects one of the three Pauli operators X, Y, Z, weighted with the given probabilities in that order.
+
+    Args:
+        probabilities (IList[float, Literal[3]]): A list of 3 probabilities corresponding to the probabilities `(p_x, p_y, p_z)` in that order.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    broadcast.single_qubit_pauli_channel(px, py, pz, ilist.IList([qubit]))
+
+
+@kernel
+def two_qubit_pauli_channel(
+    probabilities: ilist.IList[float, Literal[15]], control: Qubit, target: Qubit
+) -> None:
+    """
+    Apply a Pauli product error with weighted `probabilities` to the pair of qubits.
+
+    No error is applied with the probability `1 - sum(probabilities)`.
+
+    This will randomly select one of the pauli products
+
+    `{IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}`
+
+    weighted with the corresponding list of probabilities.
+
+    **NOTE**: The order of the given probabilities must match the order of the list of Pauli products above!
+    """
+    broadcast.two_qubit_pauli_channel(
+        probabilities, ilist.IList([control]), ilist.IList([target])
+    )
+
+
+@kernel
+def qubit_loss(p: float, qubit: Qubit) -> None:
+    """
+    Apply a qubit loss channel to the given qubit.
+
+    The qubit is lost with a probability `p`.
+
+    Args:
+        p (float): Probability of the atom being lost.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    broadcast.qubit_loss(p, ilist.IList([qubit]))
+
+
+@kernel
+def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """
+    Apply a correlated qubit loss channel to the given qubits.
+
+    All qubits are lost together with a probability `p`.
+
+    Args:
+        p (float): Probability of the qubits being lost.
+        qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied.
+    """
+    broadcast.correlated_qubit_loss(p, ilist.IList([qubits]))
+
+
+# NOTE: actual stdlib that doesn't wrap statements starts here
+
+
+@kernel
+def bit_flip(p: float, qubit: Qubit) -> None:
+    """
+    Apply a bit flip error channel to the qubit with probability `p`.
+
+    Args:
+        p (float): Probability of a bit flip error being applied.
+        qubit (Qubit): The qubit to which the noise channel is applied.
+    """
+    single_qubit_pauli_channel(p, 0, 0, qubit)

bloqade/stim/__init__.py

@@ -39,4 +39,5 @@ from ._wrappers import (
     pauli_channel1 as pauli_channel1,
     pauli_channel2 as pauli_channel2,
     observable_include as observable_include,
+    correlated_qubit_loss as correlated_qubit_loss,
 )

bloqade/stim/_wrappers.py

@@ -194,3 +194,9 @@ def z_error(p: float, targets: tuple[int, ...]) -> None: ...
 
 @wraps(noise.QubitLoss)
 def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.CorrelatedQubitLoss)
+def correlated_qubit_loss(
+    probs: tuple[float, ...], targets: tuple[int, ...]
+) -> None: ...

bloqade/stim/dialects/auxiliary/emit.py

@@ -1,7 +1,6 @@
-from kirin.emit import EmitStrFrame
 from kirin.interp import MethodTable, impl
 
-from bloqade.stim.emit.stim_str import EmitStimMain
+from bloqade.stim.emit.stim_str import EmitStimMain, EmitStimFrame
 
 from . import stmts
 from ._dialect import dialect
@@ -11,7 +10,7 @@ from ._dialect import dialect
 class EmitStimAuxMethods(MethodTable):
 
     @impl(stmts.ConstInt)
-    def const_int(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstInt):
+    def const_int(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.ConstInt):
 
         out: str = f"{stmt.value}"
 
@@ -19,7 +18,7 @@ class EmitStimAuxMethods(MethodTable):
 
     @impl(stmts.ConstFloat)
     def const_float(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstFloat
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.ConstFloat
     ):
 
         out: str = f"{stmt.value:.8f}"
@@ -28,26 +27,28 @@ class EmitStimAuxMethods(MethodTable):
 
     @impl(stmts.ConstBool)
     def const_bool(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ConstBool
+        self, emit: EmitStimMain, frame: EmitStimFrame, 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):
+    def const_str(
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.ConstBool
+    ):
 
         return (stmt.value,)
 
     @impl(stmts.Neg)
-    def neg(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Neg):
+    def neg(self, emit: EmitStimMain, frame: EmitStimFrame, 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):
+    def get_rec(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.GetRecord):
 
         id: str = frame.get(stmt.id)
         out: str = f"rec[{id}]"
@@ -55,14 +56,14 @@ class EmitStimAuxMethods(MethodTable):
         return (out,)
 
     @impl(stmts.Tick)
-    def tick(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Tick):
+    def tick(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Tick):
 
-        emit.writeln(frame, "TICK")
+        frame.write_line("TICK")
 
         return ()
 
     @impl(stmts.Detector)
-    def detector(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Detector):
+    def detector(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Detector):
 
         coords: tuple[str, ...] = frame.get_values(stmt.coord)
         targets: tuple[str, ...] = frame.get_values(stmt.targets)
@@ -70,27 +71,27 @@ class EmitStimAuxMethods(MethodTable):
         coord_str: str = ", ".join(coords)
         target_str: str = " ".join(targets)
         if len(coords):
-            emit.writeln(frame, f"DETECTOR({coord_str}) {target_str}")
+            frame.write_line(f"DETECTOR({coord_str}) {target_str}")
         else:
-            emit.writeln(frame, f"DETECTOR {target_str}")
+            frame.write_line(f"DETECTOR {target_str}")
         return ()
 
     @impl(stmts.ObservableInclude)
     def obs_include(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.ObservableInclude
+        self, emit: EmitStimMain, frame: EmitStimFrame, 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}")
+        frame.write_line(f"OBSERVABLE_INCLUDE({idx}) {target_str}")
 
         return ()
 
     @impl(stmts.NewPauliString)
     def new_paulistr(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.NewPauliString
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.NewPauliString
     ):
 
         string: tuple[str, ...] = frame.get_values(stmt.string)
@@ -105,13 +106,13 @@ class EmitStimAuxMethods(MethodTable):
 
     @impl(stmts.QubitCoordinates)
     def qubit_coordinates(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.QubitCoordinates
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.QubitCoordinates
     ):
 
         coords: tuple[str, ...] = frame.get_values(stmt.coord)
         target: str = frame.get(stmt.target)
 
         coord_str: str = ", ".join(coords)
-        emit.writeln(frame, f"QUBIT_COORDS({coord_str}) {target}")
+        frame.write_line(f"QUBIT_COORDS({coord_str}) {target}")
 
         return ()

bloqade/stim/dialects/collapse/emit_str.py

@@ -1,7 +1,6 @@
-from kirin.emit import EmitStrFrame
 from kirin.interp import MethodTable, impl
 
-from bloqade.stim.emit.stim_str import EmitStimMain
+from bloqade.stim.emit.stim_str import EmitStimMain, EmitStimFrame
 
 from . import stmts
 from ._dialect import dialect
@@ -27,13 +26,13 @@ class EmitStimCollapseMethods(MethodTable):
     @impl(stmts.MXX)
     @impl(stmts.MYY)
     @impl(stmts.MZZ)
-    def get_measure(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: Measurement):
+    def get_measure(self, emit: EmitStimMain, frame: EmitStimFrame, 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)
+        frame.write_line(out)
 
         return ()
 
@@ -46,18 +45,18 @@ class EmitStimCollapseMethods(MethodTable):
     @impl(stmts.RX)
     @impl(stmts.RY)
     @impl(stmts.RZ)
-    def get_reset(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: Reset):
+    def get_reset(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: Reset):
 
         targets: tuple[str, ...] = frame.get_values(stmt.targets)
 
         out = f"{self.reset_map[stmt.name]} " + " ".join(targets)
-        emit.writeln(frame, out)
+        frame.write_line(out)
 
         return ()
 
     @impl(stmts.PPMeasurement)
     def pp_measure(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.PPMeasurement
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.PPMeasurement
     ):
         probability: str = frame.get(stmt.p)
         targets: tuple[str, ...] = tuple(
@@ -65,6 +64,6 @@ class EmitStimCollapseMethods(MethodTable):
         )
 
         out = f"MPP({probability}) " + " ".join(targets)
-        emit.writeln(frame, out)
+        frame.write_line(out)
 
         return ()

bloqade/stim/dialects/gate/emit.py

@@ -1,7 +1,6 @@
-from kirin.emit import EmitStrFrame
 from kirin.interp import MethodTable, impl
 
-from bloqade.stim.emit.stim_str import EmitStimMain
+from bloqade.stim.emit.stim_str import EmitStimMain, EmitStimFrame
 
 from . import stmts
 from ._dialect import dialect
@@ -33,11 +32,11 @@ class EmitStimGateMethods(MethodTable):
     @impl(stmts.SqrtY)
     @impl(stmts.SqrtZ)
     def single_qubit_gate(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: SingleQubitGate
+        self, emit: EmitStimMain, frame: EmitStimFrame, 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)
+        frame.write_line(res)
 
         return ()
 
@@ -47,13 +46,13 @@ class EmitStimGateMethods(MethodTable):
 
     @impl(stmts.Swap)
     def two_qubit_gate(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: ControlledTwoQubitGate
+        self, emit: EmitStimMain, frame: EmitStimFrame, 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)
+        frame.write_line(res)
 
         return ()
 
@@ -68,19 +67,19 @@ class EmitStimGateMethods(MethodTable):
     @impl(stmts.CY)
     @impl(stmts.CZ)
     def ctrl_two_qubit_gate(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: ControlledTwoQubitGate
+        self, emit: EmitStimMain, frame: EmitStimFrame, 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)
+        frame.write_line(res)
 
         return ()
 
     @impl(stmts.SPP)
-    def spp(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.SPP):
+    def spp(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.SPP):
 
         targets: tuple[str, ...] = tuple(
             targ.upper() for targ in frame.get_values(stmt.targets)
@@ -89,6 +88,6 @@ class EmitStimGateMethods(MethodTable):
             res = "SPP_DAG " + " ".join(targets)
         else:
             res = "SPP " + " ".join(targets)
-        emit.writeln(frame, res)
+        frame.write_line(res)
 
         return ()

bloqade/stim/dialects/noise/emit.py

@@ -1,7 +1,6 @@
-from kirin.emit import EmitStrFrame
 from kirin.interp import MethodTable, impl
 
-from bloqade.stim.emit.stim_str import EmitStimMain
+from bloqade.stim.emit.stim_str import EmitStimMain, EmitStimFrame
 
 from . import stmts
 from ._dialect import dialect
@@ -24,20 +23,20 @@ class EmitStimNoiseMethods(MethodTable):
     @impl(stmts.Depolarize1)
     @impl(stmts.Depolarize2)
     def single_p_error(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.Depolarize1
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Depolarize1
     ):
 
         targets: tuple[str, ...] = frame.get_values(stmt.targets)
         p: str = frame.get(stmt.p)
         name = self.single_p_error_map[stmt.name]
         res = f"{name}({p}) " + " ".join(targets)
-        emit.writeln(frame, res)
+        frame.write_line(res)
 
         return ()
 
     @impl(stmts.PauliChannel1)
     def pauli_channel1(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.PauliChannel1
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.PauliChannel1
     ):
 
         targets: tuple[str, ...] = frame.get_values(stmt.targets)
@@ -45,13 +44,13 @@ class EmitStimNoiseMethods(MethodTable):
         py: str = frame.get(stmt.py)
         pz: str = frame.get(stmt.pz)
         res = f"PAULI_CHANNEL_1({px}, {py}, {pz}) " + " ".join(targets)
-        emit.writeln(frame, res)
+        frame.write_line(res)
 
         return ()
 
     @impl(stmts.PauliChannel2)
     def pauli_channel2(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.PauliChannel2
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.PauliChannel2
     ):
 
         targets: tuple[str, ...] = frame.get_values(stmt.targets)
@@ -61,14 +60,14 @@ class EmitStimNoiseMethods(MethodTable):
         prob_str: str = ", ".join(prob)
 
         res = f"PAULI_CHANNEL_2({prob_str}) " + " ".join(targets)
-        emit.writeln(frame, res)
+        frame.write_line(res)
 
         return ()
 
     @impl(stmts.TrivialError)
     @impl(stmts.QubitLoss)
     def non_stim_error(
-        self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.TrivialError
+        self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.TrivialError
     ):
 
         targets: tuple[str, ...] = frame.get_values(stmt.targets)
@@ -76,15 +75,16 @@ class EmitStimNoiseMethods(MethodTable):
         prob_str: str = ", ".join(prob)
 
         res = f"I_ERROR[{stmt.name}]({prob_str}) " + " ".join(targets)
-        emit.writeln(frame, res)
+        frame.write_line(res)
 
         return ()
 
     @impl(stmts.TrivialCorrelatedError)
+    @impl(stmts.CorrelatedQubitLoss)
     def non_stim_corr_error(
         self,
         emit: EmitStimMain,
-        frame: EmitStrFrame,
+        frame: EmitStimFrame,
         stmt: stmts.TrivialCorrelatedError,
     ):
 
@@ -92,7 +92,11 @@ class EmitStimNoiseMethods(MethodTable):
         prob: tuple[str, ...] = frame.get_values(stmt.probs)
         prob_str: str = ", ".join(prob)
 
-        res = f"I_ERROR[{stmt.name}:{stmt.nonce}]({prob_str}) " + " ".join(targets)
-        emit.writeln(frame, res)
+        res = (
+            f"I_ERROR[{stmt.name}:{emit.correlated_error_count}]({prob_str}) "
+            + " ".join(targets)
+        )
+        emit.correlated_error_count += 1
+        frame.write_line(res)
 
         return ()

bloqade/stim/dialects/noise/stmts.py

@@ -89,9 +89,6 @@ class NonStimError(ir.Statement):
 class NonStimCorrelatedError(ir.Statement):
     name = "NonStimCorrelatedError"
     traits = frozenset({lowering.FromPythonCall()})
-    nonce: int = (
-        info.attribute()
-    )  # Must be a unique value, otherwise stim might merge two correlated errors with equal probabilities
     probs: tuple[ir.SSAValue, ...] = info.argument(types.Float)
     targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
 
@@ -109,3 +106,8 @@ class TrivialError(NonStimError):
 @statement(dialect=dialect)
 class QubitLoss(NonStimError):
     name = "loss"
+
+
+@statement(dialect=dialect)
+class CorrelatedQubitLoss(NonStimCorrelatedError):
+    name = "correlated_loss"

bloqade/stim/emit/__init__.py

@@ -1 +1,2 @@
+from . import impls as impls
 from .stim_str import FuncEmit as FuncEmit, EmitStimMain as EmitStimMain

bloqade/stim/emit/impls.py

@@ -0,0 +1,16 @@
+from kirin.interp import MethodTable, impl
+from kirin.dialects.debug import Info, dialect
+
+from bloqade.stim.emit.stim_str import EmitStimMain, EmitStimFrame
+
+
+@dialect.register(key="emit.stim")
+class EmitStimDebugMethods(MethodTable):
+
+    @impl(Info)
+    def info(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: Info):
+
+        msg: str = frame.get(stmt.msg)
+        frame.write_line(f"# {msg}")
+
+        return ()

bloqade/stim/emit/stim_str.py

@@ -1,54 +1,71 @@
-from io import StringIO
-from typing import IO, TypeVar
-from dataclasses import field, dataclass
+import sys
+from typing import IO, Generic, TypeVar, cast
+from dataclasses import dataclass
 
 from kirin import ir, interp
-from kirin.emit import EmitStr, EmitStrFrame
 from kirin.dialects import func
+from kirin.emit.abc import EmitABC, EmitFrame
 
 IO_t = TypeVar("IO_t", bound=IO)
 
 
-def _default_dialect_group() -> ir.DialectGroup:
-    from ..groups import main
+@dataclass
+class EmitStimFrame(EmitFrame[str], Generic[IO_t]):
+    io: IO_t = cast(IO_t, sys.stdout)
+
+    def write(self, value: str) -> None:
+        self.io.write(value)
 
-    return main
+    def write_line(self, value: str) -> None:
+        self.write("    " * self._indent + value + "\n")
 
 
 @dataclass
-class EmitStimMain(EmitStr):
-    keys = ["emit.stim"]
-    dialects: ir.DialectGroup = field(default_factory=_default_dialect_group)
-    file: StringIO = field(default_factory=StringIO)
+class EmitStimMain(EmitABC[EmitStimFrame, str], Generic[IO_t]):
+    io: IO_t = cast(IO_t, sys.stdout)
+    keys = ("emit.stim",)
+    void = ""
+    correlation_identifier_offset: int = 0
 
-    def initialize(self):
+    def initialize(self) -> "EmitStimMain":
         super().initialize()
-        self.file.truncate(0)
-        self.file.seek(0)
+        self.correlated_error_count = self.correlation_identifier_offset
         return self
 
-    def eval_stmt_fallback(
-        self, frame: EmitStrFrame, stmt: ir.Statement
-    ) -> tuple[str, ...]:
-        return (stmt.name,)
+    def initialize_frame(
+        self, node: ir.Statement, *, has_parent_access: bool = False
+    ) -> EmitStimFrame:
+        return EmitStimFrame(node, self.io, has_parent_access=has_parent_access)
+
+    def frame_call(
+        self, frame: EmitStimFrame, node: ir.Statement, *args: str, **kwargs: str
+    ) -> str:
+        return f"{args[0]}({', '.join(args[1:])})"
 
-    def emit_block(self, frame: EmitStrFrame, block: ir.Block) -> str | None:
-        for stmt in block.stmts:
-            result = self.eval_stmt(frame, stmt)
-            if isinstance(result, tuple):
-                frame.set_values(stmt.results, result)
-        return None
+    def get_attribute(self, frame: EmitStimFrame, node: ir.Attribute) -> str:
+        method = self.registry.get(interp.Signature(type(node)))
+        if method is None:
+            raise ValueError(f"Method not found for node: {node}")
+        return method(self, frame, node)
 
-    def get_output(self) -> str:
-        self.file.seek(0)
-        return self.file.read()
+    def reset(self):
+        self.io.truncate(0)
+        self.io.seek(0)
+
+    def eval_fallback(self, frame: EmitStimFrame, node: ir.Statement) -> tuple:
+        return tuple("" for _ in range(len(node.results)))
 
 
 @func.dialect.register(key="emit.stim")
 class FuncEmit(interp.MethodTable):
-
     @interp.impl(func.Function)
-    def emit_func(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: func.Function):
-        _ = emit.run_ssacfg_region(frame, stmt.body, ())
-        # emit.output = "\n".join(frame.body)
+    def emit_func(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: func.Function):
+        for block in stmt.body.blocks:
+            frame.current_block = block
+            for stmt_ in block.stmts:
+                frame.current_stmt = stmt_
+                res = emit.frame_eval(frame, stmt_)
+                if isinstance(res, tuple):
+                    frame.set_values(stmt_.results, res)
+
         return ()

bloqade/stim/groups.py

@@ -1,12 +1,22 @@
 from kirin import ir
 from kirin.passes import Fold, TypeInfer
-from kirin.dialects import func, lowering
+from kirin.dialects import func, debug, ssacfg, lowering
 
 from .dialects import gate, noise, collapse, auxiliary
 
 
 @ir.dialect_group(
-    [noise, gate, auxiliary, collapse, func, lowering.func, lowering.call]
+    [
+        noise,
+        gate,
+        auxiliary,
+        collapse,
+        func,
+        lowering.func,
+        lowering.call,
+        debug,
+        ssacfg,
+    ]
 )
 def main(self):
     typeinfer_pass = TypeInfer(self)