Support Squin Depolarize to Stim DEPOLARIZE1 (#367)

Support rewriting squin depolarizing noise to stim depolarizing noise
for single qubits.

---------

Co-authored-by: John Long <[email protected]>

Author: ehua7365

_typos.toml

@@ -13,3 +13,4 @@ Braket = "Braket"
 mch = "mch"
 IY = "IY"
 ket = "ket"
+typ = "typ"

src/bloqade/stim/passes/squin_to_stim.py

@@ -8,7 +8,6 @@
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
-from kirin.analysis import const
 from kirin.ir.method import Method
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
@@ -33,18 +32,11 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
         # propagate constants
         rewrite_result = fold_pass(mt)
 
-        cp_frame, _ = const.Propagate(dialects=mt.dialects).run_analysis(mt)
-        cp_results = cp_frame.entries
-
         # Assume that address analysis and
         # wrapping has been done before this pass!
 
         # Rewrite the noise statements first.
-        rewrite_result = (
-            Walk(SquinNoiseToStim(cp_results=cp_results))
-            .rewrite(mt.code)
-            .join(rewrite_result)
-        )
+        rewrite_result = Walk(SquinNoiseToStim()).rewrite(mt.code).join(rewrite_result)
 
         # Wrap Rewrite + SquinToStim can happen w/ standard walk
 

src/bloqade/stim/rewrite/squin_noise.py

@@ -1,14 +1,14 @@
-from typing import Dict, Tuple
+from typing import Tuple
 from dataclasses import dataclass
 
 from kirin.ir import SSAValue, Statement
-from kirin.analysis import const
-from kirin.dialects import py
+from kirin.dialects import py, ilist
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.squin import op, wire, noise as squin_noise, qubit
 from bloqade.stim.dialects import noise as stim_noise
 from bloqade.stim.rewrite.util import (
+    get_const_value,
     create_wire_passthrough,
     insert_qubit_idx_after_apply,
 )
@@ -17,8 +17,6 @@
 @dataclass
 class SquinNoiseToStim(RewriteRule):
 
-    cp_results: Dict[SSAValue, const.Result]
-
     def rewrite_Statement(self, node: Statement) -> RewriteResult:
         match node:
             case qubit.Apply() | qubit.Broadcast():
@@ -67,7 +65,7 @@ def rewrite_PauliError(
         assert isinstance(squin_channel, squin_noise.stmts.PauliError)
         basis = squin_channel.basis.owner
         assert isinstance(basis, op.stmts.PauliOp)
-        p = self.cp_results.get(squin_channel.p).data
+        p = get_const_value(float, squin_channel.p)
 
         p_stmt = py.Constant(p)
         p_stmt.insert_before(stmt)
@@ -90,7 +88,7 @@ def rewrite_SingleQubitPauliChannel(
         squin_channel = stmt.operator.owner
         assert isinstance(squin_channel, squin_noise.stmts.SingleQubitPauliChannel)
 
-        params = self.cp_results.get(squin_channel.params).data
+        params = get_const_value(ilist.IList, squin_channel.params)
         new_stmts = [
             p_x := py.Constant(params[0]),
             p_y := py.Constant(params[1]),
@@ -117,7 +115,7 @@ def rewrite_TwoQubitPauliChannel(
         squin_channel = stmt.operator.owner
         assert isinstance(squin_channel, squin_noise.stmts.TwoQubitPauliChannel)
 
-        params = self.cp_results.get(squin_channel.params).data
+        params = get_const_value(ilist.IList, squin_channel.params)
         param_stmts = [py.Constant(p) for p in params]
         for param_stmt in param_stmts:
             param_stmt.insert_before(stmt)
@@ -141,3 +139,20 @@ def rewrite_TwoQubitPauliChannel(
             pzz=param_stmts[14].result,
         )
         return stim_stmt
+
+    def rewrite_Depolarize(
+        self,
+        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        qubit_idx_ssas: Tuple[SSAValue],
+    ) -> Statement:
+        """Rewrite squin.noise.Depolarize to stim.Depolarize1."""
+
+        squin_channel = stmt.operator.owner
+        assert isinstance(squin_channel, squin_noise.stmts.Depolarize)
+
+        p = get_const_value(float, squin_channel.p)
+        p_stmt = py.Constant(p)
+        p_stmt.insert_before(stmt)
+
+        stim_stmt = stim_noise.Depolarize1(targets=qubit_idx_ssas, p=p_stmt.result)
+        return stim_stmt

src/bloqade/stim/rewrite/util.py

@@ -1,4 +1,7 @@
-from kirin import ir
+from typing import TypeVar
+
+from kirin import ir, interp
+from kirin.analysis import const
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteResult
 
@@ -205,3 +208,19 @@ def is_measure_result_used(
     Check if the result of a measure statement is used in the program.
     """
     return bool(stmt.result.uses)
+
+
+T = TypeVar("T")
+
+
+def get_const_value(typ: type[T], value: ir.SSAValue) -> T:
+    if isinstance(hint := value.hints.get("const"), const.Value):
+        data = hint.data
+        if isinstance(data, typ):
+            return hint.data
+        raise interp.InterpreterError(
+            f"Expected constant value <type = {typ}>, got {data}"
+        )
+    raise interp.InterpreterError(
+        f"Expected constant value <type = {typ}>, got {value}"
+    )

test/stim/passes/stim_reference_programs/noise/apply_depolarize1.stim

@@ -0,0 +1 @@
+DEPOLARIZE1(0.01000000) 0

test/stim/passes/stim_reference_programs/noise/broadcast_depolarize1.stim

@@ -0,0 +1 @@
+DEPOLARIZE1(0.01000000) 0 1 2 3

test/stim/passes/test_squin_noise_to_stim.py

@@ -150,6 +150,34 @@ def test():
     assert codegen(test) == expected_stim_program
 
 
+def test_apply_depolarize1():
+
+    @kernel
+    def test():
+        q = qubit.new(1)
+        channel = noise.depolarize(p=0.01)
+        qubit.apply(channel, q[0])
+        return
+
+    run_address_and_stim_passes(test)
+    expected_stim_program = load_reference_program("apply_depolarize1.stim")
+    assert codegen(test) == expected_stim_program
+
+
+def test_broadcast_depolarize1():
+
+    @kernel
+    def test():
+        q = qubit.new(4)
+        channel = noise.depolarize(p=0.01)
+        qubit.broadcast(channel, q)
+        return
+
+    run_address_and_stim_passes(test)
+    expected_stim_program = load_reference_program("broadcast_depolarize1.stim")
+    assert codegen(test) == expected_stim_program
+
+
 def test_broadcast_iid_bit_flip_channel():
 
     @kernel