Adding custom rewrite for ccx gate (#245)

closes #241 

I have also added an option to `ParallelToUop` to decompose to the
rydberg gateset before just to make sure the rewrite rule is applied
properly.

---------

Co-authored-by: Xiu-zhe (Roger) Luo <[email protected]>

Author: weinbe58

src/bloqade/qasm2/passes/parallel.py

@@ -27,6 +27,7 @@
     RaiseRegisterRule,
     UOpToParallelRule,
     SimpleOptimalMergePolicy,
+    RydbergGateSetRewriteRule,
 )
 from bloqade.squin.analysis import schedule
 
@@ -135,6 +136,7 @@ def test():
     """
 
     merge_policy_type: Type[MergePolicyABC] = SimpleOptimalMergePolicy
+    rewrite_to_native_first: bool = False
     constprop: const.Propagate = field(init=False)
 
     def __post_init__(self):
@@ -147,6 +149,13 @@ def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         if not result.has_done_something:
             return result
 
+        if self.rewrite_to_native_first:
+            result = (
+                Fixpoint(Walk(RydbergGateSetRewriteRule(self.dialects)))
+                .rewrite(mt.code)
+                .join(result)
+            )
+
         frame, _ = self.constprop.run_analysis(mt)
         result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
 

src/bloqade/qasm2/rewrite/__init__.py

@@ -3,6 +3,7 @@
     GlobalToParallelRule as GlobalToParallelRule,
 )
 from .register import RaiseRegisterRule as RaiseRegisterRule
+from .native_gates import RydbergGateSetRewriteRule as RydbergGateSetRewriteRule
 from .parallel_to_uop import ParallelToUOpRule as ParallelToUOpRule
 from .uop_to_parallel import (
     MergePolicyABC as MergePolicyABC,

src/bloqade/qasm2/rewrite/native_gates.py

@@ -173,6 +173,53 @@ def rewrite_sx(self, node: uop.SX) -> abc.RewriteResult:
             cirq.XPowGate(exponent=0.5).on(self.cached_qubits[0]), node
         )
 
+    def rewrite_ccx(self, node: uop.CCX) -> abc.RewriteResult:
+        # from https://algassert.com/quirk#circuit=%7B%22cols%22:%5B%5B%22QFT3%22%5D,%5B%22inputA3%22,1,1,%22+=A3%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22%E2%80%A2%22,%22X%22%5D,%5B1,1,1,%22%E2%80%A6%22,%22%E2%80%A6%22,%22%E2%80%A6%22%5D,%5B1,1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E-%C2%BC%22%5D,%5B1,1,1,%22%E2%80%A2%22,1,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E%C2%BC%22%5D,%5B1,1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E-%C2%BC%22%5D,%5B1,1,1,%22Z%5E%C2%BC%22,%22Z%5E%C2%BC%22%5D,%5B1,1,1,1,%22H%22%5D,%5B1,1,1,%22%E2%80%A2%22,1,%22Z%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,%22X%5E-%C2%BC%22,%22X%5E%C2%BC%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,%22H%22%5D%5D%7D
+
+        # x^(1/4)
+        lam1, theta1, phi1 = map(
+            self.const_float,
+            map(around, (1.5707963267948966, 0.7853981633974483, -1.5707963267948966)),
+        )
+        lam1.insert_before(node)
+        theta1.insert_before(node)
+        phi1.insert_before(node)
+
+        lam1 = lam1.result
+        theta1 = theta1.result
+        phi1 = phi1.result
+
+        # x^(-1/4)
+        lam2, theta2, phi2 = map(
+            self.const_float,
+            map(around, (4.71238898038469, 0.7853981633974483, 1.5707963267948966)),
+        )
+        lam2.insert_before(node)
+        theta2.insert_before(node)
+        phi2.insert_before(node)
+        lam2 = lam2.result
+        theta2 = theta2.result
+        phi2 = phi2.result
+
+        uop.CZ(ctrl=node.ctrl1, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta1, phi1, lam1).insert_before(node)
+        uop.CZ(ctrl=node.ctrl1, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.T(node.ctrl1).insert_before(node)
+        uop.T(node.ctrl2).insert_before(node)
+        uop.H(node.ctrl1).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.qarg).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.ctrl1).insert_before(node)
+        uop.UGate(node.ctrl1, theta2, phi2, lam2).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.ctrl1).insert_before(node)
+        uop.H(node.ctrl1).insert_before(node)
+        node.delete()  # delete the original CCX gate
+
+        return abc.RewriteResult(has_done_something=True)
+
     def rewrite_sxdg(self, node: uop.SXdag) -> abc.RewriteResult:
         return self._rewrite_1q_gates(
             cirq.XPowGate(exponent=-0.5).on(self.cached_qubits[0]), node
@@ -394,9 +441,12 @@ def _rewrite_1q_gates(
         new_gate_stmts = self._generate_1q_gate_stmts(cirq_gate, node.qarg)
         return self._rewrite_gate_stmts(new_gate_stmts, node)
 
-    def _generate_2q_ctrl_gate_stmts(
+    def _generate_multi_ctrl_gate_stmts(
         self, cirq_gate: cirq.Operation, qubits_ssa: List[ir.SSAValue]
     ) -> list[ir.Statement]:
+        qubit_to_ssa_map = {
+            q: ssa for q, ssa in zip(self.cached_qubits[: len(qubits_ssa)], qubits_ssa)
+        }
         target_gates = self.gateset.decompose_to_target_gateset(cirq_gate, 0)
         new_stmts = []
         for new_gate in target_gates:
@@ -412,26 +462,39 @@ def _generate_2q_ctrl_gate_stmts(
                 new_stmts.append(phi2_stmt)
                 new_stmts.append(
                     uop.UGate(
-                        qarg=qubits_ssa[new_gate.qubits[0].x],
+                        qarg=qubit_to_ssa_map[new_gate.qubits[0]],
                         theta=phi0_stmt.result,
                         phi=phi1_stmt.result,
                         lam=phi2_stmt.result,
                     )
                 )
             else:
                 # 2q
-                new_stmts.append(uop.CZ(ctrl=qubits_ssa[0], qarg=qubits_ssa[1]))
+                new_stmts.append(
+                    uop.CZ(
+                        ctrl=qubit_to_ssa_map[new_gate.qubits[0]],
+                        qarg=qubit_to_ssa_map[new_gate.qubits[1]],
+                    )
+                )
 
         return new_stmts
 
     def _rewrite_2q_ctrl_gates(
         self, cirq_gate: cirq.Operation, node: uop.TwoQubitCtrlGate
     ) -> abc.RewriteResult:
-        new_gate_stmts = self._generate_2q_ctrl_gate_stmts(
+        new_gate_stmts = self._generate_multi_ctrl_gate_stmts(
             cirq_gate, [node.ctrl, node.qarg]
         )
         return self._rewrite_gate_stmts(new_gate_stmts, node)
 
+    def _rewrite_3q_ctrl_gates(
+        self, cirq_gate: cirq.Operation, node: uop.CCX
+    ) -> abc.RewriteResult:
+        new_gate_stmts = self._generate_multi_ctrl_gate_stmts(
+            cirq_gate, [node.ctrl1, node.ctrl2, node.qarg]
+        )
+        return self._rewrite_gate_stmts(new_gate_stmts, node)
+
     def _rewrite_gate_stmts(
         self, new_gate_stmts: list[ir.Statement], node: ir.Statement
     ):

test/qasm2/test_native.py

@@ -2,13 +2,15 @@
 import textwrap
 
 import cirq
+import numpy as np
 import cirq.testing
 import cirq.contrib.qasm_import as qasm_import
 import cirq.circuits.qasm_output as qasm_output
 from pytest import mark
 from kirin.rewrite import walk
 
 from bloqade import qasm2
+from bloqade.pyqrack import DynamicMemorySimulator
 from bloqade.qasm2.rewrite.native_gates import (
     RydbergGateSetRewriteRule,
     one_qubit_gate_to_u3_angles,
@@ -153,3 +155,25 @@ def kernel():
 
     # simple-stupid test to see if the rewrite injected a bunch of new lines
     assert new_qasm2.count("\n") > prog.count("\n")
+
+
+def test_ccx_rewrite():
+
+    @qasm2.extended
+    def main():
+        q = qasm2.qreg(3)
+        qasm2.ccx(q[0], q[1], q[2])
+
+        return q
+
+    main2 = main.similar()
+
+    walk.Walk(RydbergGateSetRewriteRule(main.dialects)).rewrite(main.code)
+
+    sim = DynamicMemorySimulator()
+
+    state = sim.state_vector(main)
+    state2 = sim.state_vector(main2)
+    assert (
+        np.abs(np.vdot(state, state2)) - 1 < 1e-6
+    )  # Should be close to 1 if the states are equal