Adding noise statements to squin (#211)

Some initial friction points/anticipated changes I can immediately see
in translating this to stim

- The original `PauliChannel` definition from @kaihsin can take a
variable number of float arguments, but in the stim dialect there are
explicit fields for each probability of a 1Q/2Q pauli operator being
applied. My feeling is that this should be reflected on the squin side
as just something like `1QPauliChannel` and `2QPauliChannel` with
explicit fields. I don't think we'd lose any generality here along with
the added benefit that it would be much harder for a user to misinput
information (the default definition assumes you plug in either 3 or 15
values via varargs but it seems to easy to make a mistake here)
- ~~If we DO want varargs I'm not quite sure how to nicely feed this to
the kirin `@wraps` decorator~~
- From #200 , I see it could be possible to not have to provide an
operator for `basis` and instead use a `CliffordString` (although I'd
have to restrict CliffordString even further, to just the Pauli
Operators. Would it be worth having something like a `PauliString`? Or
is that silly?)

cc: @weinbe58 @Roger-luo @kaihsin

---------

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

Author: johnzl-777

src/bloqade/squin/__init__.py

@@ -1,2 +1,2 @@
-from . import op as op, wire as wire, qubit as qubit
+from . import op as op, wire as wire, noise as noise, qubit as qubit
 from .groups import wired as wired, kernel as kernel

src/bloqade/squin/noise/__init__.py

@@ -0,0 +1,27 @@
+# Put all the proper wrappers here
+
+from kirin.lowering import wraps as _wraps
+
+from bloqade.squin.op.types import Op
+
+from . import stmts as stmts
+
+
+@_wraps(stmts.PauliError)
+def pauli_error(basis: Op, p: float) -> Op: ...
+
+
+@_wraps(stmts.PPError)
+def pp_error(op: Op, p: float) -> Op: ...
+
+
+@_wraps(stmts.Depolarize)
+def depolarize(n_qubits: int, p: float) -> Op: ...
+
+
+@_wraps(stmts.PauliChannel)
+def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
+
+
+@_wraps(stmts.QubitLoss)
+def qubit_loss(p: float) -> Op: ...

src/bloqade/squin/noise/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect(name="squin.noise")

src/bloqade/squin/noise/stmts.py

@@ -0,0 +1,59 @@
+from kirin import ir, types
+from kirin.decl import info, statement
+
+from bloqade.squin.op.types import OpType
+
+from ._dialect import dialect
+
+
+@statement
+class NoiseChannel(ir.Statement):
+    pass
+
+
+@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)
+class PPError(NoiseChannel):
+    """
+    Pauli Product Error
+    """
+
+    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
+    """
+
+    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)
+
+
+@statement(dialect=dialect)
+class QubitLoss(NoiseChannel):
+    # NOTE: qubit loss error (not supported by Stim)
+    p: ir.SSAValue = info.argument(types.Float)
+    result: ir.ResultValue = info.result(OpType)