Automatically rewrite squin noise statements before running simulator (#385)

Author: david-pl

src/bloqade/pyqrack/device.py

@@ -3,7 +3,9 @@
 
 import numpy as np
 from kirin import ir
+from kirin.passes import fold
 
+from bloqade.squin import noise as squin_noise
 from pyqrack.pauli import Pauli
 from bloqade.device import AbstractSimulatorDevice
 from bloqade.pyqrack.reg import Measurement, PyQrackQubit
@@ -16,6 +18,7 @@
     _default_pyqrack_args,
 )
 from bloqade.pyqrack.task import PyQrackSimulatorTask
+from bloqade.squin.noise.rewrite import RewriteNoiseStmts
 from bloqade.analysis.address.lattice import AnyAddress
 from bloqade.analysis.address.analysis import AddressAnalysis
 
@@ -47,14 +50,23 @@ def new_task(
         kwargs: dict[str, Any],
         memory: MemoryType,
     ) -> PyQrackSimulatorTask[Params, RetType, MemoryType]:
+
+        if squin_noise in mt.dialects:
+            # NOTE: rewrite noise statements
+            mt_ = mt.similar(mt.dialects)
+            RewriteNoiseStmts(mt_.dialects)(mt_)
+            fold.Fold(mt_.dialects)(mt_)
+        else:
+            mt_ = mt
+
         interp = PyQrackInterpreter(
-            mt.dialects,
+            mt_.dialects,
             memory=memory,
             rng_state=self.rng_state,
             loss_m_result=self.loss_m_result,
         )
         return PyQrackSimulatorTask(
-            kernel=mt, args=args, kwargs=kwargs, pyqrack_interp=interp
+            kernel=mt_, args=args, kwargs=kwargs, pyqrack_interp=interp
         )
 
     def state_vector(

test/cirq_utils/noise/test_noise_models.py

@@ -82,7 +82,7 @@ def test_simple_model(model: cirq.NoiseModel, qubits):
         assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
         assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
 
-        assert pops[0] < 0.5
-        assert pops[3] < 0.5
+        assert pops[0] < 0.5001
+        assert pops[3] < 0.5001
         assert pops[1] > 0.0
         assert pops[2] > 0.0

test/pyqrack/squin/test_noise.py

@@ -1,5 +1,7 @@
+import math
+
 from bloqade import squin
-from bloqade.pyqrack import PyQrack, PyQrackQubit
+from bloqade.pyqrack import PyQrack, PyQrackQubit, StackMemorySimulator
 from bloqade.squin.noise.stmts import NoiseChannel, StochasticUnitaryChannel
 from bloqade.squin.noise.rewrite import RewriteNoiseStmts
 
@@ -157,3 +159,53 @@ def main():
 
     target = PyQrack(1)
     target.run(main)
+
+
+def test_without_rewrite():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        x = squin.op.x()
+        squin.qubit.apply(x, q[0])
+
+        x_err = squin.noise.pauli_error(x, 0.1)
+        squin.qubit.apply(x_err, q[0])
+        return q
+
+    sim = StackMemorySimulator(min_qubits=1)
+    sim.state_vector(main)
+
+    main.print()
+
+    # make sure the method wasn't updated in-place
+    stmts = list(main.callable_region.blocks[0].stmts)
+    assert sum([isinstance(s, StochasticUnitaryChannel) for s in stmts]) == 0
+    assert sum([isinstance(s, squin.noise.stmts.PauliError) for s in stmts]) == 1
+
+
+def test_pauli_string_error():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        x = squin.op.x()
+
+        err = squin.noise.pauli_error(x, 1.0)
+        squin.qubit.apply(err, q[0])
+
+        s = squin.op.pauli_string(string="XX")
+        err2 = squin.noise.pauli_error(s, 1.0)
+        squin.qubit.apply(err2, q)
+
+    main.print()
+
+    RewriteNoiseStmts(main.dialects)(main)
+
+    main.print()
+
+    sim = StackMemorySimulator(min_qubits=2)
+
+    ket = sim.state_vector(main)
+
+    print(ket)
+
+    assert math.isclose(abs(ket[2]) ** 2, 1.0, abs_tol=1e-5)