bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/rewrite/heuristic_noise.py

@@ -0,0 +1,247 @@
+from typing import Dict, List, Tuple
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.rewrite import abc as result_abc, result
+from kirin.dialects import py, ilist
+
+from bloqade.noise import native
+from bloqade.analysis import address
+from bloqade.qasm2.dialects import uop, core, glob, parallel
+
+
+@dataclass
+class NoiseRewriteRule(result_abc.RewriteRule):
+    """
+    NOTE: This pass is not guaranteed to be supported long-term in bloqade. We will be
+    moving towards a more general approach to noise modeling in the future.
+    """
+
+    address_analysis: Dict[ir.SSAValue, address.Address]
+    gate_noise_params: native.GateNoiseParams = field(
+        default_factory=native.GateNoiseParams
+    )
+    noise_model: native.MoveNoiseModelABC = field(
+        default_factory=native.TwoRowZoneModel
+    )
+    qubit_ssa_value: Dict[int, ir.SSAValue] = field(default_factory=dict, init=False)
+
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        if isinstance(node, core.QRegNew):
+            return self.rewrite_qreg_new(node)
+        elif isinstance(node, uop.SingleQubitGate):
+            return self.rewrite_single_qubit_gate(node)
+        elif isinstance(node, uop.CZ):
+            return self.rewrite_cz_gate(node)
+        elif isinstance(node, (parallel.UGate, parallel.RZ)):
+            return self.rewrite_parallel_single_qubit_gate(node)
+        elif isinstance(node, parallel.CZ):
+            return self.rewrite_parallel_cz_gate(node)
+        elif isinstance(node, glob.UGate):
+            return self.rewrite_global_single_qubit_gate(node)
+        else:
+            return result.RewriteResult()
+
+    def rewrite_qreg_new(self, node: core.QRegNew):
+
+        addr = self.address_analysis[node.result]
+        if not isinstance(addr, address.AddressReg):
+            return result.RewriteResult()
+
+        has_done_something = False
+        for idx_val, qid in enumerate(addr.data):
+            if qid not in self.qubit_ssa_value:
+                has_done_something = True
+                idx = py.constant.Constant(value=idx_val)
+                qubit = core.QRegGet(node.result, idx=idx.result)
+                self.qubit_ssa_value[qid] = qubit.result
+                qubit.insert_after(node)
+                idx.insert_after(node)
+
+        return result.RewriteResult(has_done_something=has_done_something)
+
+    def insert_single_qubit_noise(
+        self,
+        node: ir.Statement,
+        qargs: ir.SSAValue,
+        probs: Tuple[float, float, float, float],
+    ):
+        native.PauliChannel(qargs, px=probs[0], py=probs[1], pz=probs[2]).insert_before(
+            node
+        )
+        native.AtomLossChannel(qargs, prob=probs[3]).insert_before(node)
+
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_single_qubit_gate(self, node: uop.SingleQubitGate):
+        probs = (
+            self.gate_noise_params.local_px,
+            self.gate_noise_params.local_py,
+            self.gate_noise_params.local_pz,
+            self.gate_noise_params.local_loss_prob,
+        )
+        (qargs := ilist.New(values=(node.qarg,))).insert_before(node)
+        return self.insert_single_qubit_noise(node, qargs.result, probs)
+
+    def rewrite_global_single_qubit_gate(self, node: glob.UGate):
+        addrs = self.address_analysis[node.registers]
+        if not isinstance(addrs, address.AddressTuple):
+            return result.RewriteResult()
+
+        qargs = []
+
+        for addr in addrs.data:
+            if not isinstance(addr, address.AddressReg):
+                return result.RewriteResult()
+
+            for qid in addr.data:
+                qargs.append(self.qubit_ssa_value[qid])
+
+        probs = (
+            self.gate_noise_params.global_px,
+            self.gate_noise_params.global_py,
+            self.gate_noise_params.global_pz,
+            self.gate_noise_params.global_loss_prob,
+        )
+        (qargs := ilist.New(values=tuple(qargs))).insert_before(node)
+        return self.insert_single_qubit_noise(node, qargs.result, probs)
+
+    def rewrite_parallel_single_qubit_gate(self, node: parallel.RZ | parallel.UGate):
+        addrs = self.address_analysis[node.qargs]
+        if not isinstance(addrs, address.AddressTuple):
+            return result.RewriteResult()
+
+        if not all(isinstance(addr, address.AddressQubit) for addr in addrs.data):
+            return result.RewriteResult()
+
+        probs = (
+            self.gate_noise_params.local_px,
+            self.gate_noise_params.local_py,
+            self.gate_noise_params.local_pz,
+            self.gate_noise_params.local_loss_prob,
+        )
+        assert isinstance(node.qargs, ir.ResultValue)
+        assert isinstance(node.qargs.stmt, ilist.New)
+        return self.insert_single_qubit_noise(node, node.qargs, probs)
+
+    def move_noise_stmts(
+        self,
+        errors: Dict[Tuple[float, float, float, float], List[int]],
+    ) -> list[ir.Statement]:
+
+        nodes = []
+
+        for probs, qubits in errors.items():
+            if len(qubits) == 0:
+                continue
+
+            nodes.append(
+                qargs := ilist.New(tuple(self.qubit_ssa_value[q] for q in qubits))
+            )
+            nodes.append(native.AtomLossChannel(qargs.result, prob=probs[3]))
+            nodes.append(
+                native.PauliChannel(qargs.result, px=probs[0], py=probs[1], pz=probs[2])
+            )
+
+        return nodes
+
+    def cz_gate_noise(
+        self,
+        ctrls: ir.SSAValue,
+        qargs: ir.SSAValue,
+    ) -> list[ir.Statement]:
+        return [
+            native.CZPauliChannel(
+                ctrls,
+                qargs,
+                px_ctrl=self.gate_noise_params.cz_paired_gate_px,
+                py_ctrl=self.gate_noise_params.cz_paired_gate_py,
+                pz_ctrl=self.gate_noise_params.cz_paired_gate_pz,
+                px_qarg=self.gate_noise_params.cz_paired_gate_px,
+                py_qarg=self.gate_noise_params.cz_paired_gate_py,
+                pz_qarg=self.gate_noise_params.cz_paired_gate_pz,
+                paired=True,
+            ),
+            native.CZPauliChannel(
+                ctrls,
+                qargs,
+                px_ctrl=self.gate_noise_params.cz_unpaired_gate_px,
+                py_ctrl=self.gate_noise_params.cz_unpaired_gate_py,
+                pz_ctrl=self.gate_noise_params.cz_unpaired_gate_pz,
+                px_qarg=self.gate_noise_params.cz_unpaired_gate_px,
+                py_qarg=self.gate_noise_params.cz_unpaired_gate_py,
+                pz_qarg=self.gate_noise_params.cz_unpaired_gate_pz,
+                paired=False,
+            ),
+            native.AtomLossChannel(
+                ctrls, prob=self.gate_noise_params.cz_gate_loss_prob
+            ),
+            native.AtomLossChannel(
+                qargs, prob=self.gate_noise_params.cz_gate_loss_prob
+            ),
+        ]
+
+    def rewrite_cz_gate(self, node: uop.CZ):
+
+        has_done_something = False
+
+        qarg_addr = self.address_analysis[node.qarg]
+        ctrl_addr = self.address_analysis[node.ctrl]
+
+        (ctrls := ilist.New([node.ctrl])).insert_before(node)
+        (qargs := ilist.New([node.qarg])).insert_before(node)
+
+        if isinstance(qarg_addr, address.AddressQubit) and isinstance(
+            ctrl_addr, address.AddressQubit
+        ):
+            other_qubits = sorted(
+                set(self.qubit_ssa_value.keys()) - {ctrl_addr.data, qarg_addr.data}
+            )
+            errors = self.noise_model.parallel_cz_errors(
+                [ctrl_addr.data], [qarg_addr.data], other_qubits
+            )
+
+            move_noise_nodes = self.move_noise_stmts(errors)
+
+            for new_node in move_noise_nodes:
+                new_node.insert_before(node)
+                has_done_something = True
+
+        gate_noise_nodes = self.cz_gate_noise(ctrls.result, qargs.result)
+
+        for new_node in gate_noise_nodes:
+            new_node.insert_before(node)
+            has_done_something = True
+
+        return result.RewriteResult(has_done_something=has_done_something)
+
+    def rewrite_parallel_cz_gate(self, node: parallel.CZ):
+        ctrls = self.address_analysis[node.ctrls]
+        qargs = self.address_analysis[node.qargs]
+
+        has_done_something = False
+        if (
+            isinstance(ctrls, address.AddressTuple)
+            and all(isinstance(addr, address.AddressQubit) for addr in ctrls.data)
+            and isinstance(qargs, address.AddressTuple)
+            and all(isinstance(addr, address.AddressQubit) for addr in qargs.data)
+        ):
+            ctrl_qubits = list(map(lambda addr: addr.data, ctrls.data))
+            qarg_qubits = list(map(lambda addr: addr.data, qargs.data))
+            rest = sorted(
+                set(self.qubit_ssa_value.keys()) - set(ctrl_qubits + qarg_qubits)
+            )
+            errors = self.noise_model.parallel_cz_errors(ctrl_qubits, qarg_qubits, rest)
+            move_noise_nodes = self.move_noise_stmts(errors)
+
+            for new_node in move_noise_nodes:
+                new_node.insert_before(node)
+                has_done_something = True
+
+        gate_noise_nodes = self.cz_gate_noise(node.ctrls, node.qargs)
+
+        for new_node in gate_noise_nodes:
+            new_node.insert_before(node)
+            has_done_something = True
+
+        return result.RewriteResult(has_done_something=has_done_something)

bloqade/qasm2/rewrite/native_gates.py

@@ -0,0 +1,447 @@
+import math
+from typing import List, Optional
+from functools import cached_property
+from dataclasses import field, dataclass
+
+import cirq
+import numpy as np
+import cirq.transformers
+import cirq.contrib.qasm_import
+import cirq.transformers.target_gatesets
+import cirq.transformers.target_gatesets.compilation_target_gateset
+from kirin import ir
+from kirin.rewrite import abc, result
+from kirin.dialects import py
+from cirq.circuits.qasm_output import QasmUGate
+from cirq.transformers.target_gatesets.compilation_target_gateset import (
+    CompilationTargetGateset,
+)
+
+from bloqade.qasm2.dialects import uop, expr
+
+
+# rydeberg gate sets
+class RydbergTargetGateset(cirq.CZTargetGateset):
+    def __init__(self, *, cnz_max_size: int = 2, atol: float = 1e-8):
+        additional = [cirq.Z.controlled(cn) for cn in range(2, cnz_max_size)]
+        super().__init__(atol=atol, additional_gates=additional)
+        self.cnz_max_size = cnz_max_size
+
+    @property
+    def num_qubits(self) -> int:
+        return max(2, self.cnz_max_size)
+
+
+# decompose the CU by defining a custom Cirq Gate with the qelib1 definition
+# Need to be careful about the fact that U(theta, phi, lambda) in standard QASM2
+# and its variants
+class CU(cirq.Gate):
+    def __init__(self, theta, phi, lam, gamma):
+        super(CU, self)
+        self.theta = theta
+        self.phi = phi
+        self.lam = lam
+        self.gamma = gamma
+
+    def _num_qubits_(self):
+        return 2
+
+    def _decompose_(self, qubits):
+        ctrl, target = qubits
+        # taken from qelib1 definition
+        # p(gamma) c;
+        yield QasmUGate(0, 0, self.gamma / math.pi)(ctrl)
+        # p((lambda+phi/2)) c;
+        yield QasmUGate(0, 0, ((self.lam + self.phi) / 2) / math.pi)(ctrl)
+        # p((lambda-phi/2)) t;
+        yield QasmUGate(0, 0, ((self.lam - self.phi) / 2) / math.pi)(target)
+        # cx c,t
+        yield cirq.CX(ctrl, target)
+        # u(-theta/2, 0, -(phi+lambda/2)) t;
+        yield QasmUGate(
+            (-self.theta / 2) / math.pi, 0, (-(self.phi + self.lam) / 2) / math.pi
+        )(target)
+        # cx c,t
+        yield cirq.CX(ctrl, target)
+        # u(theta/2, phi, 0) t;
+        yield QasmUGate((self.theta / 2) / math.pi, self.phi / math.pi, 0)(target)
+
+    def _circuit_diagram_info_(self, args):
+        return "*", "CU"
+
+
+def around(val):
+    return float(np.around(val, 14))
+
+
+def one_qubit_gate_to_u3_angles(op: cirq.Operation) -> tuple[float, float, float]:
+    lam, theta, phi = (  # Z angle, Y angle, then Z angle
+        cirq.deconstruct_single_qubit_matrix_into_angles(cirq.unitary(op))
+    )
+    return tuple(map(around, (theta, phi, lam)))
+
+
+@dataclass
+class RydbergGateSetRewriteRule(abc.RewriteRule):
+    # NOTE
+    # 1. this can only rewrite qasm2.main and qasm2.gate!
+    dialect_group: ir.DialectGroup
+    gateset: CompilationTargetGateset = field(default_factory=RydbergTargetGateset)
+
+    @cached_property
+    def cached_qubits(self) -> tuple[cirq.LineQubit, ...]:
+
+        # qasm2 stmts only have up to 3 qubits gates, so we cached only 3.
+        return tuple(cirq.LineQubit(i) for i in range(3))
+
+    @cached_property
+    def const_float_type(self):
+        if expr.dialect in self.dialect_group.data:
+            return expr.ConstFloat
+        else:
+            return py.constant.Constant
+
+    def const_float(self, value: float):
+        return self.const_float_type(value=value)
+
+    @cached_property
+    def const_pi(self):
+        if expr in self.dialect_group.data:
+            return expr.ConstPI()
+        else:
+            return py.constant.Constant(value=math.pi)
+
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        # only deal with uop
+        if type(node) in uop.dialect.stmts:
+            return getattr(self, f"rewrite_{node.name}")(node)
+
+        return result.RewriteResult()
+
+    def rewrite_barrier(self, node: uop.Barrier) -> result.RewriteResult:
+        return result.RewriteResult()
+
+    def rewrite_cz(self, node: uop.CZ) -> result.RewriteResult:
+        return result.RewriteResult()
+
+    def rewrite_CX(self, node: uop.CX) -> result.RewriteResult:
+        return self._rewrite_2q_ctrl_gates(
+            cirq.CX(self.cached_qubits[0], self.cached_qubits[1]), node
+        )
+
+    def rewrite_cy(self, node: uop.CY) -> result.RewriteResult:
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ControlledGate(cirq.Y, 1)(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+
+    def rewrite_U(self, node: uop.UGate) -> result.RewriteResult:
+        return result.RewriteResult()
+
+    def rewrite_id(self, node: uop.Id) -> result.RewriteResult:
+        node.delete()  # just delete the identity gate
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_h(self, node: uop.H) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.H(self.cached_qubits[0]), node)
+
+    def rewrite_x(self, node: uop.X) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.X(self.cached_qubits[0]), node)
+
+    def rewrite_y(self, node: uop.Y) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.Y(self.cached_qubits[0]), node)
+
+    def rewrite_z(self, node: uop.Z) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.Z(self.cached_qubits[0]), node)
+
+    def rewrite_s(self, node: uop.S) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.S(self.cached_qubits[0]), node)
+
+    def rewrite_sdg(self, node: uop.Sdag) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.S(self.cached_qubits[0]) ** -1, node)
+
+    def rewrite_t(self, node: uop.T) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.T(self.cached_qubits[0]), node)
+
+    def rewrite_tdg(self, node: uop.Tdag) -> result.RewriteResult:
+        return self._rewrite_1q_gates(cirq.T(self.cached_qubits[0]) ** -1, node)
+
+    def rewrite_sx(self, node: uop.SX) -> result.RewriteResult:
+        return self._rewrite_1q_gates(
+            cirq.XPowGate(exponent=0.5).on(self.cached_qubits[0]), node
+        )
+
+    def rewrite_sxdg(self, node: uop.SXdag) -> result.RewriteResult:
+        return self._rewrite_1q_gates(
+            cirq.XPowGate(exponent=-0.5).on(self.cached_qubits[0]), node
+        )
+
+    def rewrite_u1(self, node: uop.U1) -> result.RewriteResult:
+        theta = node.lam
+        (phi := self.const_float(value=0.0)).insert_before(node)
+        node.replace_by(
+            uop.UGate(qarg=node.qarg, theta=phi.result, phi=phi.result, lam=theta)
+        )
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_u2(self, node: uop.U2) -> result.RewriteResult:
+        phi = node.phi
+        lam = node.lam
+        (theta := self.const_float(value=math.pi / 2)).insert_before(node)
+        node.replace_by(uop.UGate(qarg=node.qarg, theta=theta.result, phi=phi, lam=lam))
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_rx(self, node: uop.RX) -> result.RewriteResult:
+        theta = node.theta
+        (phi := self.const_float(value=math.pi / 2)).insert_before(node)
+        (lam := self.const_float(value=-math.pi / 2)).insert_before(node)
+        node.replace_by(
+            uop.UGate(qarg=node.qarg, theta=theta, phi=phi.result, lam=lam.result)
+        )
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_ry(self, node: uop.RY) -> result.RewriteResult:
+        theta = node.theta
+        (phi := self.const_float(value=0.0)).insert_before(node)
+        node.replace_by(
+            uop.UGate(qarg=node.qarg, theta=theta, phi=phi.result, lam=phi.result)
+        )
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_rz(self, node: uop.RZ) -> result.RewriteResult:
+        theta = node.theta
+        (phi := self.const_float(value=0.0)).insert_before(node)
+        node.replace_by(
+            uop.UGate(qarg=node.qarg, theta=phi.result, phi=phi.result, lam=theta)
+        )
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_crx(self, node: uop.CRX) -> result.RewriteResult:
+        lam = self._get_const_value(node.lam)
+
+        if lam is None:
+            return result.RewriteResult()
+
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ControlledGate(cirq.Rx(rads=lam), 1).on(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+
+    def rewrite_cry(self, node: uop.CRY) -> result.RewriteResult:
+        lam = self._get_const_value(node.lam)
+
+        if lam is None:
+            return result.RewriteResult()
+
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ControlledGate(cirq.Ry(rads=lam), 1).on(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+
+    def rewrite_crz(self, node: uop.CRZ) -> result.RewriteResult:
+        lam = self._get_const_value(node.lam)
+
+        if lam is None:
+            return result.RewriteResult()
+
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ControlledGate(cirq.Rz(rads=lam), 1).on(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+
+    def rewrite_cu1(self, node: uop.CU1) -> result.RewriteResult:
+
+        lam = self._get_const_value(node.lam)
+
+        if lam is None:
+            return result.RewriteResult()
+
+        # cirq.ControlledGate(u3(0, 0, lambda))
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ControlledGate(QasmUGate(0, 0, lam / math.pi)).on(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+        pass
+
+    def rewrite_cu3(self, node: uop.CU3) -> result.RewriteResult:
+
+        theta = self._get_const_value(node.theta)
+        lam = self._get_const_value(node.lam)
+        phi = self._get_const_value(node.phi)
+
+        if not all((theta, phi, lam)):
+            return result.RewriteResult()
+
+        # cirq.ControlledGate(u3(theta, lambda phi))
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ControlledGate(
+                QasmUGate(theta / math.pi, phi / math.pi, lam / math.pi)
+            ).on(self.cached_qubits[0], self.cached_qubits[1]),
+            node,
+        )
+
+    def rewrite_cu(self, node: uop.CU) -> result.RewriteResult:
+
+        gamma = self._get_const_value(node.gamma)
+        theta = self._get_const_value(node.theta)
+        lam = self._get_const_value(node.lam)
+        phi = self._get_const_value(node.phi)
+
+        # need to create custom 2q gate, then feed that into rewrite_2q
+
+        return self._rewrite_2q_ctrl_gates(
+            CU(theta, phi, lam, gamma).on(self.cached_qubits[0], self.cached_qubits[1]),
+            node,
+        )
+
+    def rewrite_rxx(self, node: uop.RXX) -> result.RewriteResult:
+
+        theta = self._get_const_value(node.theta)
+
+        if theta is None:
+            return result.RewriteResult()
+
+        # even though the XX gate is not controlled,
+        # the end U + CZ decomposition that happens internally means
+        return self._rewrite_2q_ctrl_gates(
+            cirq.XXPowGate(exponent=theta / math.pi).on(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+
+    def rewrite_rzz(self, node: uop.RZZ) -> result.RewriteResult:
+        theta = self._get_const_value(node.theta)
+
+        if theta is None:
+            return result.RewriteResult()
+
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ZZPowGate(exponent=theta / math.pi).on(
+                self.cached_qubits[0], self.cached_qubits[1]
+            ),
+            node,
+        )
+
+        """
+        return self._rewrite_2q_ctrl_gates(
+            cirq.ZZPowGate(exponent = theta).on(self.cached_qubits[0], self.cached_qubits[1])
+            ,node
+        )
+        """
+
+    def rewrite_swap(self, node: uop.Swap):
+        return self._rewrite_2q_ctrl_gates(
+            cirq.SWAP(self.cached_qubits[0], self.cached_qubits[1]), node
+        )
+
+    def _get_const_value(self, ssa: ir.SSAValue) -> Optional[float | int]:
+        if not isinstance(ssa, ir.ResultValue):
+            return None
+
+        match ssa.owner:
+            case expr.ConstFloat(value=value):
+                return value
+            case expr.ConstInt(value=value):
+                return value
+            case py.constant.Constant(value=float() as value) | py.constant.Constant(
+                value=int() as value
+            ):
+                return value
+            case expr.ConstPI():
+                return math.pi
+            case _:
+                return None
+
+    def _generate_1q_gate_stmts(self, cirq_gate: cirq.Operation, qarg: ir.SSAValue):
+        target_gates = self.gateset.decompose_to_target_gateset(cirq_gate, 0)
+
+        if isinstance(target_gates, cirq.GateOperation):
+            target_gates = [target_gates]
+
+        new_stmts = []
+        for new_gate in target_gates:
+            theta, phi, lam = one_qubit_gate_to_u3_angles(new_gate)
+            theta_stmt = self.const_float(value=theta)
+            phi_stmt = self.const_float(value=phi)
+            lam_stmt = self.const_float(value=lam)
+
+            new_stmts.append(theta_stmt)
+            new_stmts.append(phi_stmt)
+            new_stmts.append(lam_stmt)
+            new_stmts.append(
+                uop.UGate(
+                    qarg=qarg,
+                    theta=theta_stmt.result,
+                    phi=phi_stmt.result,
+                    lam=lam_stmt.result,
+                )
+            )
+        return new_stmts
+
+    def _rewrite_1q_gates(
+        self, cirq_gate: cirq.Operation, node: uop.SingleQubitGate
+    ) -> result.RewriteResult:
+        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(
+        self, cirq_gate: cirq.Operation, qubits_ssa: List[ir.SSAValue]
+    ) -> list[ir.Statement]:
+        target_gates = self.gateset.decompose_to_target_gateset(cirq_gate, 0)
+        new_stmts = []
+        for new_gate in target_gates:
+            if len(new_gate.qubits) == 1:
+                # 1q
+                phi0, phi1, phi2 = one_qubit_gate_to_u3_angles(new_gate)
+                phi0_stmt = self.const_float(value=phi0)
+                phi1_stmt = self.const_float(value=phi1)
+                phi2_stmt = self.const_float(value=phi2)
+
+                new_stmts.append(phi0_stmt)
+                new_stmts.append(phi1_stmt)
+                new_stmts.append(phi2_stmt)
+                new_stmts.append(
+                    uop.UGate(
+                        qarg=qubits_ssa[new_gate.qubits[0].x],
+                        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]))
+
+        return new_stmts
+
+    def _rewrite_2q_ctrl_gates(
+        self, cirq_gate: cirq.Operation, node: uop.TwoQubitCtrlGate
+    ) -> result.RewriteResult:
+        new_gate_stmts = self._generate_2q_ctrl_gate_stmts(
+            cirq_gate, [node.ctrl, 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
+    ):
+
+        node.replace_by(new_gate_stmts[0])
+        node = new_gate_stmts[0]
+
+        for stmt in new_gate_stmts[1:]:
+            stmt.insert_after(node)
+            node = stmt
+
+        return result.RewriteResult(has_done_something=True)