implement translation of squin qubit loss to stim dialect (#322)

Translate the squin qubit loss statement into special stim form.

The necessity of this was brought to my attention by @ChenZhao44

---------

Co-authored-by: Kai-Hsin Wu <[email protected]>

Author: johnzl-777

src/bloqade/stim/__init__.py

@@ -31,6 +31,7 @@
     z_error as z_error,
     detector as detector,
     identity as identity,
+    qubit_loss as qubit_loss,
     depolarize1 as depolarize1,
     depolarize2 as depolarize2,
     pauli_string as pauli_string,

src/bloqade/stim/_wrappers.py

@@ -190,3 +190,7 @@ def y_error(p: float, targets: tuple[int, ...]) -> None: ...
 
 @wraps(noise.ZError)
 def z_error(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.QubitLoss)
+def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ...

src/bloqade/stim/dialects/noise/emit.py

@@ -66,6 +66,7 @@ def pauli_channel2(
         return ()
 
     @impl(stmts.TrivialError)
+    @impl(stmts.QubitLoss)
     def non_stim_error(
         self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.TrivialError
     ):

src/bloqade/stim/dialects/noise/stmts.py

@@ -104,3 +104,8 @@ class TrivialCorrelatedError(NonStimCorrelatedError):
 @statement(dialect=dialect)
 class TrivialError(NonStimError):
     name = "TRIV_ERROR"
+
+
+@statement(dialect=dialect)
+class QubitLoss(NonStimError):
+    name = "loss"

src/bloqade/stim/rewrite/qubit_to_stim.py

@@ -1,11 +1,12 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, qubit
+from bloqade.squin import op, noise, qubit
 from bloqade.squin.rewrite import AddressAttribute
 from bloqade.stim.rewrite.util import (
-    SQUIN_STIM_GATE_MAPPING,
+    SQUIN_STIM_OP_MAPPING,
     rewrite_Control,
+    rewrite_QubitLoss,
     insert_qubit_idx_from_address,
 )
 
@@ -29,14 +30,18 @@ def rewrite_Apply_and_Broadcast(
 
         # this is an SSAValue, need it to be the actual operator
         applied_op = stmt.operator.owner
+
+        if isinstance(applied_op, noise.stmts.QubitLoss):
+            return rewrite_QubitLoss(stmt)
+
         assert isinstance(applied_op, op.stmts.Operator)
 
         if isinstance(applied_op, op.stmts.Control):
             return rewrite_Control(stmt)
 
         # need to handle Control through separate means
         # but we can handle X, Y, Z, H, and S here just fine
-        stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
+        stim_1q_op = SQUIN_STIM_OP_MAPPING.get(type(applied_op))
         if stim_1q_op is None:
             return RewriteResult()
 
@@ -45,7 +50,6 @@ def rewrite_Apply_and_Broadcast(
         if address_attr is None:
             return RewriteResult()
 
-        # sometimes you can get a whole AddressReg...
         assert isinstance(address_attr, AddressAttribute)
         qubit_idx_ssas = insert_qubit_idx_from_address(
             address=address_attr, stmt_to_insert_before=stmt

src/bloqade/stim/rewrite/util.py

@@ -2,19 +2,20 @@
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteResult
 
-from bloqade.squin import op, wire, qubit
+from bloqade.squin import op, wire, noise as squin_noise, qubit
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.stim.dialects import gate, collapse
+from bloqade.stim.dialects import gate, noise as stim_noise, collapse
 from bloqade.analysis.address import AddressReg, AddressWire, AddressQubit, AddressTuple
 
-SQUIN_STIM_GATE_MAPPING = {
+SQUIN_STIM_OP_MAPPING = {
     op.stmts.X: gate.X,
     op.stmts.Y: gate.Y,
     op.stmts.Z: gate.Z,
     op.stmts.H: gate.H,
     op.stmts.S: gate.S,
     op.stmts.Identity: gate.Identity,
     op.stmts.Reset: collapse.RZ,
+    squin_noise.stmts.QubitLoss: stim_noise.QubitLoss,
 }
 
 # Squin allows creation of control gates where the gate can be any operator,
@@ -149,18 +150,52 @@ def rewrite_Control(
     stim_stmt = stim_gate(controls=ctrl_qubits, targets=target_qubits)
 
     if isinstance(stmt_with_ctrl, (wire.Apply, wire.Broadcast)):
-        # have to "reroute" the input of these statements to directly plug in
-        # to subsequent statements, remove dependency on the current statement
-        for input_wire, output_wire in zip(
-            stmt_with_ctrl.inputs, stmt_with_ctrl.results
-        ):
-            output_wire.replace_by(input_wire)
+        create_wire_passthrough(stmt_with_ctrl)
 
     stmt_with_ctrl.replace_by(stim_stmt)
 
     return RewriteResult(has_done_something=True)
 
 
+def rewrite_QubitLoss(
+    stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+) -> RewriteResult:
+    """
+    Rewrite QubitLoss statements to Stim's TrivialError.
+    """
+
+    squin_loss_op = stmt.operator.owner
+    assert isinstance(squin_loss_op, squin_noise.stmts.QubitLoss)
+
+    qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt)
+    if qubit_idx_ssas is None:
+        return RewriteResult()
+
+    stim_loss_stmt = stim_noise.QubitLoss(
+        targets=qubit_idx_ssas,
+        probs=(squin_loss_op.p,),
+    )
+
+    if isinstance(stmt, (wire.Apply, wire.Broadcast)):
+        create_wire_passthrough(stmt)
+
+    stmt.replace_by(stim_loss_stmt)
+    # NoiseChannels are not pure,
+    # need to manually delete because
+    # DCE won't touch them
+    stmt.operator.owner.delete()
+
+    return RewriteResult(has_done_something=True)
+
+
+def create_wire_passthrough(stmt: wire.Apply | wire.Broadcast) -> None:
+
+    for input_wire, output_wire in zip(stmt.inputs, stmt.results):
+        # have to "reroute" the input of these statements to directly plug in
+        # to subsequent statements, remove dependency on the current statement
+        output_wire.replace_by(input_wire)
+
+
 def is_measure_result_used(
     stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure,
 ) -> bool:

src/bloqade/stim/rewrite/wire_to_stim.py

@@ -1,10 +1,11 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, wire
+from bloqade.squin import op, wire, noise
 from bloqade.stim.rewrite.util import (
-    SQUIN_STIM_GATE_MAPPING,
+    SQUIN_STIM_OP_MAPPING,
     rewrite_Control,
+    rewrite_QubitLoss,
     insert_qubit_idx_from_wire_ssa,
 )
 
@@ -24,12 +25,16 @@ def rewrite_Apply_and_Broadcast(
 
         # this is an SSAValue, need it to be the actual operator
         applied_op = stmt.operator.owner
+
+        if isinstance(applied_op, noise.stmts.QubitLoss):
+            return rewrite_QubitLoss(stmt)
+
         assert isinstance(applied_op, op.stmts.Operator)
 
         if isinstance(applied_op, op.stmts.Control):
             return rewrite_Control(stmt)
 
-        stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
+        stim_1q_op = SQUIN_STIM_OP_MAPPING.get(type(applied_op))
         if stim_1q_op is None:
             return RewriteResult()
 

test/stim/dialects/stim/test_stim_circuits.py

@@ -86,6 +86,13 @@ def test_pauli_error():
     interp.run(test_pauli_error, args=())
     print(interp.get_output())
 
+    @stim.main
+    def test_qubit_loss():
+        stim.qubit_loss(probs=(0.1, 0.2), targets=(0, 1, 2))
+
+    interp.run(test_qubit_loss, args=())
+    assert interp.get_output() == "\nI_ERROR[loss](0.10000000, 0.20000000) 0 1 2"
+
 
 def test_collapse():
     @stim.main

test/stim/parse/test_parse_custom.py

@@ -28,3 +28,15 @@ def test_parse_trivial_correlated():
     # test roundtrip
     out = codegen(mt)
     assert out.strip() == "I_ERROR[TRIV_CORR_ERROR:3](0.20000000, 0.30000000) 5 0 1 2"
+
+
+def test_qubit_loss():
+    mt = loads(
+        "I_ERROR[loss](0.1, 0.2) 0 1", nonstim_noise_ops={"loss": noise.QubitLoss}
+    )
+
+    mt.print()
+
+    # test roundtrip
+    out = codegen(mt)
+    assert out.strip() == "I_ERROR[loss](0.10000000, 0.20000000) 0 1"

test/stim/passes/stim_reference_programs/qubit/qubit.txt renamed to test/stim/passes/stim_reference_programs/qubit/qubit.stim

@@ -0,0 +1,5 @@
+
+H 0 1 2 3 4
+I_ERROR[loss](0.10000000) 3
+I_ERROR[loss](0.05000000) 0 1 2 3 4
+MZ(0.00000000) 0 1 2 3 4

test/stim/passes/stim_reference_programs/qubit/qubit_broadcast.txt renamed to test/stim/passes/stim_reference_programs/qubit/qubit_broadcast.stim

@@ -0,0 +1,8 @@
+
+I_ERROR[loss](0.10000000) 0 1 2 3 4
+I_ERROR[loss](0.90000000) 0
+MZ(0.00000000) 0
+MZ(0.00000000) 1
+MZ(0.00000000) 2
+MZ(0.00000000) 3
+MZ(0.00000000) 4

test/stim/passes/stim_reference_programs/qubit/qubit_loss.stim

@@ -5,7 +5,7 @@
 from kirin.dialects import py
 
 from bloqade import squin
-from bloqade.squin import op, qubit, kernel
+from bloqade.squin import op, noise, qubit, kernel
 from bloqade.stim.emit import EmitStimMain
 from bloqade.stim.passes import SquinToStim
 from bloqade.squin.rewrite import WrapAddressAnalysis
@@ -37,7 +37,7 @@ def load_reference_program(filename):
         return f.read()
 
 
-def run_address_and_stim_passes(test):
+def run_address_and_stim_passes(test: ir.Method):
     addr_frame, _ = AddressAnalysis(test.dialects).run_analysis(test)
     Walk(WrapAddressAnalysis(address_analysis=addr_frame.entries)).rewrite(test.code)
     SquinToStim(test.dialects)(test)
@@ -57,7 +57,7 @@ def test():
         return
 
     run_address_and_stim_passes(test)
-    base_stim_prog = load_reference_program("qubit.txt")
+    base_stim_prog = load_reference_program("qubit.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()
 
@@ -74,7 +74,7 @@ def test():
         return
 
     run_address_and_stim_passes(test)
-    base_stim_prog = load_reference_program("qubit_reset.txt")
+    base_stim_prog = load_reference_program("qubit_reset.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()
 
@@ -91,6 +91,26 @@ def test():
         return
 
     run_address_and_stim_passes(test)
-    base_stim_prog = load_reference_program("qubit_broadcast.txt")
+    base_stim_prog = load_reference_program("qubit_broadcast.stim")
+
+    assert codegen(test) == base_stim_prog.rstrip()
+
+
+def test_qubit_loss():
+    @kernel
+    def test():
+        n_qubits = 5
+        ql = qubit.new(n_qubits)
+        # apply Hadamard to all qubits
+        squin.qubit.broadcast(op.h(), ql)
+        # apply and broadcast qubit loss
+        squin.qubit.apply(noise.qubit_loss(0.1), ql[3])
+        squin.qubit.broadcast(noise.qubit_loss(0.05), ql)
+        # measure out
+        squin.qubit.measure(ql)
+        return
+
+    run_address_and_stim_passes(test)
+    base_stim_prog = load_reference_program("qubit_loss.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()