Fix missing adjoint problem (#562)

Author: johnzl-777

pyproject.toml

@@ -13,7 +13,7 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
     "scipy>=1.13.1",
-    "kirin-toolchain~=0.17.23",
+    "kirin-toolchain~=0.17.26",
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",

src/bloqade/stim/rewrite/qubit_to_stim.py

@@ -3,7 +3,6 @@
 
 from bloqade.squin import op, noise, qubit
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.stim.dialects import gate
 from bloqade.stim.rewrite.util import (
     SQUIN_STIM_OP_MAPPING,
     rewrite_Control,
@@ -40,20 +39,24 @@ def rewrite_Apply_and_Broadcast(
 
         assert isinstance(applied_op, op.stmts.Operator)
 
+        # Handle controlled gates with a separate procedure
         if isinstance(applied_op, op.stmts.Control):
             return rewrite_Control(stmt)
 
-        # need to handle Control through separate means
-
         # check if its adjoint, assume its canonicalized so no nested adjoints.
         is_conj = False
         if isinstance(applied_op, op.stmts.Adjoint):
-            if not applied_op.is_unitary:
+            # By default the Adjoint has is_unitary = False, so we need to check
+            # the inner applied operator to make sure its not just unitary,
+            # but something that has an equivalent stim representation with *_DAG format.
+            if isinstance(
+                applied_op.op.owner, (op.stmts.SqrtX, op.stmts.SqrtY, op.stmts.S)
+            ):
+                is_conj = True
+                applied_op = applied_op.op.owner
+            else:
                 return RewriteResult()
 
-            is_conj = True
-            applied_op = applied_op.op.owner
-
         stim_1q_op = SQUIN_STIM_OP_MAPPING.get(type(applied_op))
         if stim_1q_op is None:
             return RewriteResult()
@@ -71,13 +74,14 @@ def rewrite_Apply_and_Broadcast(
         if qubit_idx_ssas is None:
             return RewriteResult()
 
-        if isinstance(stim_1q_op, gate.stmts.Gate):
+        # At this point, we know for certain stim_1q_op must be SQRT_X, SQRT_Y, or S
+        # and has the option to set the dagger attribute. If is_conj is false,
+        # the rewrite would have terminated early so we know anything else has to be
+        # a non 1Q gate operation.
+        if is_conj:
             stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas), dagger=is_conj)
         else:
             stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
         stmt.replace_by(stim_1q_stmt)
 
         return RewriteResult(has_done_something=True)
-
-
-# put rewrites for measure statements in separate rule, then just have to dispatch

test/stim/passes/stim_reference_programs/qubit/adjoint_rewrite.stim

@@ -0,0 +1,4 @@
+
+SQRT_X_DAG 0
+SQRT_Y_DAG 0
+S_DAG 0

test/stim/passes/stim_reference_programs/qubit/u3_to_clifford.stim

@@ -1,3 +1,12 @@
 
 H 0
+S 0
+SQRT_Y 0
+S 0
+S_DAG 0
+SQRT_Y 0
+S 0
+S 0
+SQRT_Y 0
+S_DAG 0
 MZ(0.00000000) 0

test/stim/passes/test_squin_qubit_to_stim.py

@@ -110,6 +110,26 @@ def test():
     assert codegen(test) == base_stim_prog.rstrip()
 
 
+def test_adjoint_rewrite():
+
+    @squin.kernel
+    def test():
+        q = qubit.new(1)
+        sqrt_x_dag = op.adjoint(op.sqrt_x())
+        sqrt_y_dag = op.adjoint(op.sqrt_y())
+        sqrt_s_dag = op.adjoint(op.s())
+        qubit.apply(sqrt_x_dag, q[0])
+        qubit.apply(sqrt_y_dag, q[0])
+        qubit.apply(sqrt_s_dag, q[0])
+        return
+
+    SquinToStimPass(test.dialects)(test)
+
+    base_stim_prog = load_reference_program("adjoint_rewrite.stim")
+
+    assert codegen(test) == base_stim_prog.rstrip()
+
+
 def test_u3_to_clifford():
 
     @kernel
@@ -118,12 +138,24 @@ def test():
         q = qubit.new(n_qubits)
         # apply U3 rotation that can be translated to a Clifford gate
         squin.qubit.apply(op.u(0.25 * math.tau, 0.0 * math.tau, 0.5 * math.tau), q[0])
+        # S @ SQRT_Y @ S = Z @ SQRT_X
+        squin.qubit.apply(
+            op.u(-0.25 * math.tau, -0.25 * math.tau, -0.25 * math.tau), q[0]
+        )
+        # S @ SQRT_Y @ S_DAG = SQRT_X_DAG
+        squin.qubit.apply(
+            op.u(-0.25 * math.tau, -0.25 * math.tau, 0.25 * math.tau), q[0]
+        )
+        # S_DAG @ SQRT_Y @ S = SQRT_X
+        squin.qubit.apply(
+            op.u(-0.25 * math.tau, 0.25 * math.tau, -0.25 * math.tau), q[0]
+        )
+
         # measure out
         squin.qubit.measure(q)
         return
 
     SquinToStimPass(test.dialects)(test)
-
     base_stim_prog = load_reference_program("u3_to_clifford.stim")
 
     assert codegen(test) == base_stim_prog.rstrip()