Adding native noise dialect to the extended dialect by default (#269)

This should enable you guys to program noise inside a kernel by default.

cc: @stefan-o @MilanKornjacaQ @Roger-luo

---------

Co-authored-by: David Plankensteiner <[email protected]>

Author: weinbe58

src/bloqade/noise/native/_wrappers.py

@@ -8,15 +8,28 @@
 
 
 @wraps(native.AtomLossChannel)
-def atom_loss_channel(
-    qargs: ilist.IList[Qubit, Any] | list, *, prob: float
-) -> None: ...
+def atom_loss_channel(qargs: ilist.IList[Qubit, Any] | list, *, prob: float) -> None:
+    """Apply an atom loss channel to a list of qubits.
+
+    Args:
+        qargs (ilist.IList[Qubit, Any] | list): List of qubits to apply the noise to.
+        prob (float): The loss probability.
+    """
+    ...
 
 
 @wraps(native.PauliChannel)
 def pauli_channel(
     qargs: ilist.IList[Qubit, Any] | list, *, px: float, py: float, pz: float
-) -> None: ...
+) -> None:
+    """Apply a Pauli channel to a list of qubits.
+
+    Args:
+        qargs (ilist.IList[Qubit, Any] | list): List of qubits to apply the noise to.
+        px (float): Probability of X error.
+        py (float): Probability of Y error.
+        pz (float): Probability of Z error.
+    """
 
 
 @wraps(native.CZPauliChannel)
@@ -31,4 +44,20 @@ def cz_pauli_channel(
     py_qarg: float,
     pz_qarg: float,
     paired: bool,
-) -> None: ...
+) -> None:
+    """Insert noise for a CZ gate with a Pauli channel on qubits.
+
+    Args:
+        ctrls: List of control qubits.
+        qarg2: List of target qubits.
+        px_ctrl: Probability of X error on control qubits.
+        py_ctrl: Probability of Y error on control qubits.
+        pz_ctrl: Probability of Z error on control qubits.
+        px_qarg: Probability of X error on target qubits.
+        py_qarg: Probability of Y error on target qubits.
+        pz_qarg: Probability of Z error on target qubits.
+        paired: If True, the noise is applied to both control and target qubits
+            are not lost otherwise skip this error. If False Apply the noise on
+            the whatever qubit is not lost.
+    """
+    ...

src/bloqade/qasm2/groups.py

@@ -2,6 +2,7 @@
 from kirin.prelude import structural_no_opt
 from kirin.dialects import scf, func, ilist, lowering
 
+from bloqade.noise import native
 from bloqade.qasm2.dialects import (
     uop,
     core,
@@ -90,6 +91,7 @@ def run_pass(
             noise,
             parallel,
             core,
+            native,
         ]
     )
 )

test/analysis/fidelity/test_fidelity.py

@@ -7,8 +7,6 @@
 from bloqade.analysis.fidelity import FidelityAnalysis
 from bloqade.qasm2.passes.noise import NoisePass
 
-noise_main = qasm2.extended.add(native.dialect)
-
 
 class NoiseTestModel(native.MoveNoiseModelABC):
     def parallel_cz_errors(self, ctrls, qargs, rest):
@@ -17,7 +15,7 @@ def parallel_cz_errors(self, ctrls, qargs, rest):
 
 def test_basic_noise():
 
-    @noise_main
+    @qasm2.extended
     def main():
         q = qasm2.qreg(2)
         qasm2.x(q[0])
@@ -65,7 +63,7 @@ def main():
 
 
 def test_c_noise():
-    @noise_main
+    @qasm2.extended
     def main():
         q = qasm2.qreg(2)
         qasm2.cz(q[0], q[1])
@@ -118,7 +116,7 @@ def main():
 @pytest.mark.xfail
 def test_if():
 
-    @noise_main
+    @qasm2.extended
     def main():
         q = qasm2.qreg(1)
         c = qasm2.creg(1)
@@ -129,7 +127,7 @@ def main():
 
         return c
 
-    @noise_main
+    @qasm2.extended
     def main_if():
         q = qasm2.qreg(1)
         c = qasm2.creg(1)
@@ -179,7 +177,7 @@ def main_if():
 @pytest.mark.xfail
 def test_for():
 
-    @noise_main
+    @qasm2.extended
     def main():
         q = qasm2.qreg(1)
         c = qasm2.creg(1)
@@ -195,7 +193,7 @@ def main():
 
         return c
 
-    @noise_main
+    @qasm2.extended
     def main_for():
         q = qasm2.qreg(1)
         c = qasm2.creg(1)

test/pyqrack/runtime/noise/native/test_loss.py

@@ -9,8 +9,6 @@
 from bloqade.pyqrack import PyQrackQubit, PyQrackInterpreter, reg
 from bloqade.pyqrack.base import MockMemory
 
-simulation = qasm2.extended.add(native)
-
 
 def run_mock(program: ir.Method, rng_state: Mock | None = None):
     PyQrackInterpreter(
@@ -22,7 +20,7 @@ def run_mock(program: ir.Method, rng_state: Mock | None = None):
 
 def test_atom_loss():
 
-    @simulation
+    @qasm2.extended
     def test_atom_loss(c: qasm2.CReg):
         q = qasm2.qreg(2)
         native.atom_loss_channel([q[0]], prob=0.1)
@@ -37,7 +35,7 @@ def test_atom_loss(c: qasm2.CReg):
     memory = MockMemory()
 
     result: ilist.IList[PyQrackQubit, Literal[2]] = PyQrackInterpreter(
-        simulation, memory=memory, rng_state=rng_state
+        qasm2.extended, memory=memory, rng_state=rng_state
     ).run(test_atom_loss, (input,))
 
     assert result[0].state is reg.QubitState.Lost

test/pyqrack/runtime/noise/native/test_pauli.py

@@ -7,8 +7,6 @@
 from bloqade.noise import native
 from bloqade.pyqrack.base import MockMemory, PyQrackInterpreter
 
-simulation = qasm2.extended.add(native)
-
 
 def run_mock(program: ir.Method, rng_state: Mock | None = None):
     PyQrackInterpreter(
@@ -19,7 +17,7 @@ def run_mock(program: ir.Method, rng_state: Mock | None = None):
 
 
 def test_pauli_channel():
-    @simulation
+    @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
         native.pauli_channel(
@@ -44,7 +42,7 @@ def test_atom_loss():
 
 @pytest.mark.xfail
 def test_pauli_probs_check():
-    @simulation
+    @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
         native.pauli_channel(
@@ -60,7 +58,7 @@ def test_atom_loss():
 
 
 def test_cz_pauli_channel_false():
-    @simulation
+    @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
         native.atom_loss_channel([q[0]], prob=0.4)
@@ -101,7 +99,7 @@ def test_atom_loss():
 
 
 def test_cz_pauli_channel_true():
-    @simulation
+    @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
         native.atom_loss_channel([q[0]], prob=0.4)