bloqade-circuit 0.2.3__py3-none-any.whl → 0.3.0__py3-none-any.whl

bloqade/qasm2/emit/main.py

@@ -5,8 +5,11 @@ from kirin.dialects import cf, scf, func
 from kirin.ir.dialect import Dialect as Dialect
 
 from bloqade.qasm2.parse import ast
+from bloqade.qasm2.dialects.uop import SingleQubitGate, TwoQubitCtrlGate
+from bloqade.qasm2.dialects.expr import GateFunction
 
 from .base import EmitQASM2Base, EmitQASM2Frame
+from ..dialects.core.stmts import Reset, Measure
 
 
 @dataclass
@@ -94,6 +97,24 @@ class Scf(interp.MethodTable):
 
         cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
 
+        # NOTE: we need exactly one of those in the then body in order to emit valid QASM2
+        AllowedThenType = SingleQubitGate | TwoQubitCtrlGate | Measure | Reset
+
+        then_stmts = stmt.then_body.blocks[0].stmts
+        uop_stmts = 0
+        for s in then_stmts:
+            if isinstance(s, AllowedThenType):
+                uop_stmts += 1
+                continue
+
+            if isinstance(s, func.Invoke):
+                uop_stmts += isinstance(s.callee.code, GateFunction)
+
+        if uop_stmts != 1:
+            raise interp.InterpreterError(
+                "Cannot lower if-statement: QASM2 only allows exactly one quantum operation in the body."
+            )
+
         with emit.new_frame(stmt) as then_frame:
             then_frame.entries.update(frame.entries)
             emit.emit_block(then_frame, stmt.then_body.blocks[0])

bloqade/qasm2/emit/target.py

@@ -11,6 +11,7 @@ from bloqade.qasm2.passes.glob import GlobalToParallel
 from bloqade.qasm2.passes.py2qasm import Py2QASM
 from bloqade.qasm2.passes.parallel import ParallelToUOp
 
+from . import impls as impls  # register the tables
 from .gate import EmitQASM2Gate
 from .main import EmitQASM2Main
 
@@ -27,6 +28,8 @@ class QASM2:
         allow_parallel: bool = False,
         allow_global: bool = False,
         custom_gate: bool = True,
+        unroll_ifs: bool = True,
+        allow_noise: bool = True,
     ) -> None:
         """Initialize the QASM2 target.
 
@@ -43,13 +46,18 @@ class QASM2:
             qelib1 (bool):
                 Include the `include "qelib1.inc"` line in the resulting QASM2 AST that's
                 submitted to qBraid. Defaults to `True`.
+
             custom_gate (bool):
                 Include the custom gate definitions in the resulting QASM2 AST. Defaults to `True`. If `False`, all the qasm2.gate will be inlined.
 
+            unroll_ifs (bool):
+                Unrolls if statements with multiple qasm2 statements in the body in order to produce valid qasm2 output, which only allows a single
+                operation in an if body. Defaults to `True`.
+
 
 
         """
-        from bloqade import qasm2
+        from bloqade import noise, qasm2
 
         self.main_target = qasm2.main
         self.gate_target = qasm2.gate
@@ -58,6 +66,7 @@ class QASM2:
         self.custom_gate = custom_gate
         self.allow_parallel = allow_parallel
         self.allow_global = allow_global
+        self.unroll_ifs = unroll_ifs
 
         if allow_parallel:
             self.main_target = self.main_target.add(qasm2.dialects.parallel)
@@ -67,7 +76,11 @@ class QASM2:
             self.main_target = self.main_target.add(qasm2.dialects.glob)
             self.gate_target = self.gate_target.add(qasm2.dialects.glob)
 
-        if allow_global or allow_parallel:
+        if allow_noise:
+            self.main_target = self.main_target.add(noise.native)
+            self.gate_target = self.gate_target.add(noise.native)
+
+        if allow_global or allow_parallel or allow_noise:
             self.main_target = self.main_target.add(ilist)
             self.gate_target = self.gate_target.add(ilist)
 
@@ -87,9 +100,11 @@ class QASM2:
 
         # make a cloned instance of kernel
         entry = entry.similar()
-        QASM2Fold(entry.dialects, inline_gate_subroutine=not self.custom_gate).fixpoint(
-            entry
-        )
+        QASM2Fold(
+            entry.dialects,
+            inline_gate_subroutine=not self.custom_gate,
+            unroll_ifs=self.unroll_ifs,
+        ).fixpoint(entry)
 
         if not self.allow_global:
             # rewrite global to parallel

bloqade/qasm2/groups.py

@@ -2,6 +2,7 @@ from kirin import ir, passes
 from kirin.prelude import structural_no_opt
 from kirin.dialects import scf, func, ilist, lowering
 
+from bloqade.noise import native
 from bloqade.qasm2.dialects import (
     uop,
     core,
@@ -90,6 +91,7 @@ def main(self):
             noise,
             parallel,
             core,
+            native,
         ]
     )
 )

bloqade/qasm2/parse/__init__.py

@@ -19,19 +19,22 @@ def loadfile(file: str | pathlib.Path):
         return loads(f.read())
 
 
-def pprint(node: ast.Node, *, console: Console | None = None):
+def pprint(node: ast.Node, *, console: Console | None = None, no_color: bool = False):
     if console:
-        return Printer(console).visit(node)
+        printer = Printer(console)
     else:
-        Printer().visit(node)
+        printer = Printer()
+    printer.console.no_color = no_color
+    printer.visit(node)
 
 
-def spprint(node: ast.Node, *, console: Console | None = None):
+def spprint(node: ast.Node, *, console: Console | None = None, no_color: bool = False):
     if console:
         printer = Printer(console)
     else:
         printer = Printer()
 
+    printer.console.no_color = no_color
     with printer.string_io() as stream:
         printer.visit(node)
         return stream.getvalue()

bloqade/qasm2/parse/lowering.py

@@ -1,5 +1,3 @@
-import os
-import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -19,131 +17,39 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     hint_show_lineno: bool = field(default=True, kw_only=True)
     stacktrace: bool = field(default=True, kw_only=True)
 
-    def loads(
+    def run(
         self,
-        source: str,
-        kernel_name: str,
+        stmt: ast.Node,
         *,
-        returns: str | None = None,
+        source: str | None = None,
         globals: dict[str, Any] | None = None,
         file: str | None = None,
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
-    ) -> ir.Method:
-        from ..parse import loads
-
-        # TODO: add source info
-        stmt = loads(source)
+    ) -> ir.Region:
 
-        state = lowering.State(
-            self,
+        frame = self.get_frame(
+            stmt,
+            source=source,
+            globals=globals,
             file=file,
             lineno_offset=lineno_offset,
             col_offset=col_offset,
         )
-        with state.frame(
-            [stmt],
-            globals=globals,
-            finalize_next=False,
-        ) as frame:
-            try:
-                self.visit(state, stmt)
-                # append return statement with the return values
-                if returns is not None:
-                    return_value = frame.get(returns)
-                    if return_value is None:
-                        raise lowering.BuildError(f"Cannot find return value {returns}")
-                else:
-                    return_value = func.ConstantNone()
-                    frame.push(return_value)
-
-                return_node = frame.push(func.Return(value_or_stmt=return_value))
-
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-
-        code = func.Function(
-            sym_name=kernel_name,
-            signature=func.Signature((), return_node.value.type),
-            body=region,
-        )
 
-        return ir.Method(
-            mod=None,
-            py_func=None,
-            sym_name=kernel_name,
-            arg_names=[],
-            dialects=self.dialects,
-            code=code,
-        )
+        return frame.curr_region
 
-    def loadfile(
-        self,
-        file: str | pathlib.Path,
-        *,
-        kernel_name: str | None = None,
-        returns: str | None = None,
-        globals: dict[str, Any] | None = None,
-        lineno_offset: int = 0,
-        col_offset: int = 0,
-        compactify: bool = True,
-    ) -> ir.Method:
-        if isinstance(file, str):
-            file = pathlib.Path(*os.path.split(file))
-
-        if not file.is_file() or not file.name.endswith(".qasm"):
-            raise ValueError("File must be a .qasm file")
-
-        kernel_name = (
-            file.name.replace(".qasm", "") if kernel_name is None else kernel_name
-        )
-
-        with file.open("r") as f:
-            source = f.read()
-
-        return self.loads(
-            source,
-            kernel_name,
-            returns=returns,
-            globals=globals,
-            file=str(file),
-            lineno_offset=lineno_offset,
-            col_offset=col_offset,
-            compactify=compactify,
-        )
-
-    def run(
+    def get_frame(
         self,
         stmt: ast.Node,
-        *,
         source: str | None = None,
         globals: dict[str, Any] | None = None,
         file: str | None = None,
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
-    ) -> ir.Region:
+    ) -> lowering.Frame:
         # TODO: add source info
         state = lowering.State(
             self,
@@ -154,32 +60,16 @@ class QASM2(lowering.LoweringABC[ast.Node]):
         with state.frame(
             [stmt],
             globals=globals,
+            finalize_next=False,
         ) as frame:
-            try:
-                self.visit(state, stmt)
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-        return region
+            self.visit(state, stmt)
+
+            if compactify:
+                from kirin.rewrite import Walk, CFGCompactify
+
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
+
+            return frame
 
     def visit(self, state: lowering.State[ast.Node], node: ast.Node) -> lowering.Result:
         name = node.__class__.__name__

bloqade/qasm2/parse/qasm2.lark

@@ -9,7 +9,7 @@ version: INT "." INT
 // stmts
 include: "include" STRING ";"
 ifstmt: "if" "(" expr "==" expr ")" ifbody
-ifbody: qop | "{" qop* "}" // allow multiple qops
+ifbody: qop
 opaque: "opaque" IDENTIFIER ["(" [params] ")"] qubits ";"
 barrier: "barrier" qubits ";"
 qreg: "qreg" IDENTIFIER "[" INT "]" ";"

bloqade/qasm2/passes/__init__.py

@@ -3,3 +3,4 @@ from .noise import NoisePass as NoisePass
 from .py2qasm import Py2QASM as Py2QASM
 from .qasm2py import QASM2Py as QASM2Py
 from .parallel import UOpToParallel as UOpToParallel
+from .unroll_if import UnrollIfs as UnrollIfs

bloqade/qasm2/passes/fold.py

@@ -23,6 +23,8 @@ from kirin.rewrite.abc import RewriteResult
 
 from bloqade.qasm2.dialects import expr
 
+from .unroll_if import UnrollIfs
+
 
 @dataclass
 class QASM2Fold(Pass):
@@ -30,6 +32,7 @@ class QASM2Fold(Pass):
 
     constprop: const.Propagate = field(init=False)
     inline_gate_subroutine: bool = True
+    unroll_ifs: bool = True
 
     def __post_init__(self):
         self.constprop = const.Propagate(self.dialects)
@@ -61,6 +64,9 @@ class QASM2Fold(Pass):
             .join(result)
         )
 
+        if self.unroll_ifs:
+            UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
+
         # run typeinfer again after unroll etc. because we now insert
         # a lot of new nodes, which might have more precise types
         self.typeinfer.unsafe_run(mt)

bloqade/qasm2/passes/noise.py

@@ -62,13 +62,32 @@ class NoisePass(Pass):
     def __post_init__(self):
         self.address_analysis = address.AddressAnalysis(self.dialects)
 
+    def get_qubit_values(self, mt: ir.Method):
+        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_values = {}
+        # Traverse statements in block order to fine the first SSA value for each qubit
+        for block in mt.callable_region.blocks:
+            for stmt in block.stmts:
+                if len(stmt.results) != 1:
+                    continue
+
+                addr = frame.entries.get(result := stmt.results[0])
+                if (
+                    isinstance(addr, address.AddressQubit)
+                    and (index := addr.data) not in qubit_ssa_values
+                ):
+                    qubit_ssa_values[index] = result
+
+        return qubit_ssa_values, frame.entries
+
     def unsafe_run(self, mt: ir.Method):
         result = LiftQubits(self.dialects).unsafe_run(mt)
-        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_value, address_analysis = self.get_qubit_values(mt)
         result = (
             Walk(
                 NoiseRewriteRule(
-                    address_analysis=frame.entries,
+                    qubit_ssa_value=qubit_ssa_value,
+                    address_analysis=address_analysis,
                     noise_model=self.noise_model,
                     gate_noise_params=self.gate_noise_params,
                 ),
@@ -77,5 +96,6 @@ class NoisePass(Pass):
             .rewrite(mt.code)
             .join(result)
         )
+
         result = Fixpoint(Walk(DeadCodeElimination())).rewrite(mt.code).join(result)
         return result

bloqade/qasm2/passes/parallel.py

@@ -27,6 +27,7 @@ from bloqade.qasm2.rewrite import (
     RaiseRegisterRule,
     UOpToParallelRule,
     SimpleOptimalMergePolicy,
+    RydbergGateSetRewriteRule,
 )
 from bloqade.squin.analysis import schedule
 
@@ -135,6 +136,7 @@ class UOpToParallel(Pass):
     """
 
     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 @@ class UOpToParallel(Pass):
         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)
 

bloqade/qasm2/passes/unroll_if.py

@@ -0,0 +1,25 @@
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    CommonSubexpressionElimination,
+)
+
+from ..rewrite.split_ifs import LiftThenBody, SplitIfStmts
+
+
+class UnrollIfs(Pass):
+    """This pass lifts statements that are not UOP out of the if body and then splits whatever is left into multiple if statements so you obtain valid QASM2"""
+
+    def unsafe_run(self, mt: ir.Method):
+        result = Walk(LiftThenBody()).rewrite(mt.code)
+        result = Walk(SplitIfStmts()).rewrite(mt.code).join(result)
+        result = (
+            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        return result

bloqade/qasm2/rewrite/__init__.py

@@ -3,6 +3,7 @@ from .glob import (
     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,

bloqade/qasm2/rewrite/desugar.py

@@ -5,16 +5,17 @@ from kirin.passes import Pass
 from kirin.rewrite import abc, walk
 from kirin.dialects import py
 
+from bloqade.qasm2 import types
 from bloqade.qasm2.dialects import core
 
 
 class IndexingDesugarRule(abc.RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, py.indexing.GetItem):
-            if node.obj.type.is_subseteq(core.QRegType):
+            if node.obj.type.is_subseteq(types.QRegType):
                 node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
-            elif node.obj.type.is_subseteq(core.CRegType):
+            elif node.obj.type.is_subseteq(types.CRegType):
                 node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
 

bloqade/qasm2/rewrite/heuristic_noise.py

@@ -18,6 +18,7 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
     """
 
     address_analysis: Dict[ir.SSAValue, address.Address]
+    qubit_ssa_value: Dict[int, ir.SSAValue]
     gate_noise_params: native.GateNoiseParams = field(
         default_factory=native.GateNoiseParams
     )
@@ -25,15 +26,6 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
         default_factory=native.TwoRowZoneModel
     )
 
-    def __post_init__(self):
-        self.qubit_ssa_value: Dict[int, ir.SSAValue] = {}
-        for ssa_value, addr in self.address_analysis.items():
-            if (
-                isinstance(addr, address.AddressQubit)
-                and ssa_value not in self.qubit_ssa_value
-            ):
-                self.qubit_ssa_value[addr.data] = ssa_value
-
     def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
         if isinstance(node, uop.SingleQubitGate):
             return self.rewrite_single_qubit_gate(node)

bloqade/qasm2/rewrite/native_gates.py

@@ -173,6 +173,53 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
             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 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
         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,7 +462,7 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
                 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,
@@ -420,18 +470,31 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
                 )
             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
     ):