Move circuit field from EmitCirqFrame to EmitCirq

Author: zhenrongliew

src/bloqade/cirq_utils/emit/base.py

@@ -6,7 +6,7 @@
 from kirin import ir, types, interp
 from kirin.emit import EmitABC, EmitFrame
 from kirin.interp import MethodTable, impl
-from kirin.dialects import py, func
+from kirin.dialects import py, func, ilist
 from typing_extensions import Self
 
 from bloqade.squin import kernel
@@ -149,17 +149,23 @@ def main():
     new_func = func.Function(
         sym_name=sym_name, body=callable_region, signature=new_signature
     )
-    mt_ = ir.Method(None, None, sym_name, [], mt.dialects, new_func)
+    # mt_ = ir.Method(None, None, sym_name, [], mt.dialects, new_func)
+    mt_ = ir.Method(
+        dialects=mt.dialects,
+        code=new_func,
+        sym_name=sym_name,
+    )
 
     AggressiveUnroll(mt_.dialects).fixpoint(mt_)
-    return emitter.run(mt_, args=())
+    emitter.initialize()
+    emitter.run(mt_)
+    return emitter.circuit
 
 
 @dataclass
 class EmitCirqFrame(EmitFrame):
     qubit_index: int = 0
     qubits: Sequence[cirq.Qid] | None = None
-    circuit: cirq.Circuit = field(default_factory=cirq.Circuit)
 
 
 def _default_kernel():
@@ -172,19 +178,20 @@ class EmitCirq(EmitABC[EmitCirqFrame, cirq.Circuit]):
     dialects: ir.DialectGroup = field(default_factory=_default_kernel)
     void = cirq.Circuit()
     qubits: Sequence[cirq.Qid] | None = None
+    circuit: cirq.Circuit = field(default_factory=cirq.Circuit)
 
     def initialize(self) -> Self:
         return super().initialize()
 
     def initialize_frame(
-        self, code: ir.Statement, *, has_parent_access: bool = False
+        self, node: ir.Statement, *, has_parent_access: bool = False
     ) -> EmitCirqFrame:
         return EmitCirqFrame(
-            code, has_parent_access=has_parent_access, qubits=self.qubits
+            node, has_parent_access=has_parent_access, qubits=self.qubits
         )
 
     def run_method(self, method: ir.Method, args: tuple[cirq.Circuit, ...]):
-        return self.run_callable(method.code, args)
+        return self.call(method, *args)
 
     def run_callable_region(
         self,
@@ -198,7 +205,7 @@ def run_callable_region(
             # NOTE: skip self arg
             frame.set_values(block_args[1:], args)
 
-        results = self.eval_stmt(frame, code)
+        results = self.frame_eval(frame, code)
         if isinstance(results, tuple):
             if len(results) == 0:
                 return self.void
@@ -208,20 +215,32 @@ def run_callable_region(
 
     def emit_block(self, frame: EmitCirqFrame, block: ir.Block) -> cirq.Circuit:
         for stmt in block.stmts:
-            result = self.eval_stmt(frame, stmt)
+            result = self.frame_eval(frame, stmt)
             if isinstance(result, tuple):
                 frame.set_values(stmt.results, result)
 
-        return frame.circuit
+        return self.circuit
+
+    def reset(self):
+        pass
 
 
 @func.dialect.register(key="emit.cirq")
 class __FuncEmit(MethodTable):
 
     @impl(func.Function)
     def emit_func(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Function):
-        emit.run_ssacfg_region(frame, stmt.body, ())
-        return (frame.circuit,)
+        for block in stmt.body.blocks:
+            frame.current_block = block
+            for s in block.stmts:
+                frame.current_stmt = s
+                stmt_results = emit.frame_eval(frame, s)
+                if isinstance(stmt_results, tuple):
+                    if len(stmt_results) != 0:
+                        frame.set_values(s.results, stmt_results)
+                    continue
+
+        return (emit.circuit,)
 
     @impl(func.Invoke)
     def emit_invoke(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Invoke):
@@ -235,6 +254,12 @@ def return_(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Return):
         # NOTE: should only be hit if ignore_returns == True
         return ()
 
+    @impl(func.ConstantNone)
+    def emit_constant_none(
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.ConstantNone
+    ):
+        return ()
+
 
 @py.indexing.dialect.register(key="emit.cirq")
 class __Concrete(interp.MethodTable):
@@ -243,3 +268,19 @@ class __Concrete(interp.MethodTable):
     def getindex(self, interp, frame: interp.Frame, stmt: py.indexing.GetItem):
         # NOTE: no support for indexing into single statements in cirq
         return ()
+
+    @interp.impl(py.Constant)
+    def emit_constant(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: py.Constant):
+        return (stmt.value.data,)  # pyright: ignore[reportAttributeAccessIssue]
+
+
+@ilist.dialect.register(key="emit.cirq")
+class __IList(interp.MethodTable):
+    @interp.impl(ilist.New)
+    def new_ilist(
+        self,
+        emit: EmitCirq,
+        frame: interp.Frame,
+        stmt: ilist.New,
+    ):
+        return (ilist.IList(data=frame.get_values(stmt.values)),)

src/bloqade/cirq_utils/emit/gate.py

@@ -20,7 +20,7 @@ def hermitian(
     ):
         qubits = frame.get(stmt.qubits)
         cirq_op = getattr(cirq, stmt.name.upper())
-        frame.circuit.append(cirq_op.on_each(qubits))
+        emit.circuit.append(cirq_op.on_each(qubits))
         return ()
 
     @impl(gate.stmts.S)
@@ -36,7 +36,7 @@ def unitary(
         if stmt.adjoint:
             cirq_op = cirq_op ** (-1)
 
-        frame.circuit.append(cirq_op.on_each(qubits))
+        emit.circuit.append(cirq_op.on_each(qubits))
         return ()
 
     @impl(gate.stmts.SqrtX)
@@ -58,7 +58,7 @@ def sqrt(
         else:
             cirq_op = cirq.YPowGate(exponent=exponent)
 
-        frame.circuit.append(cirq_op.on_each(qubits))
+        emit.circuit.append(cirq_op.on_each(qubits))
         return ()
 
     @impl(gate.stmts.CX)
@@ -71,7 +71,7 @@ def control(
         targets = frame.get(stmt.targets)
         cirq_op = getattr(cirq, stmt.name.upper())
         cirq_qubits = [(ctrl, target) for ctrl, target in zip(controls, targets)]
-        frame.circuit.append(cirq_op.on_each(cirq_qubits))
+        emit.circuit.append(cirq_op.on_each(cirq_qubits))
         return ()
 
     @impl(gate.stmts.Rx)
@@ -84,7 +84,7 @@ def rot(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.RotationGat
         angle = turns * 2 * math.pi
         cirq_op = getattr(cirq, stmt.name.title())(rads=angle)
 
-        frame.circuit.append(cirq_op.on_each(qubits))
+        emit.circuit.append(cirq_op.on_each(qubits))
         return ()
 
     @impl(gate.stmts.U3)
@@ -95,10 +95,10 @@ def u3(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.U3):
         phi = frame.get(stmt.phi) * 2 * math.pi
         lam = frame.get(stmt.lam) * 2 * math.pi
 
-        frame.circuit.append(cirq.Rz(rads=lam).on_each(*qubits))
+        emit.circuit.append(cirq.Rz(rads=lam).on_each(*qubits))
 
-        frame.circuit.append(cirq.Ry(rads=theta).on_each(*qubits))
+        emit.circuit.append(cirq.Ry(rads=theta).on_each(*qubits))
 
-        frame.circuit.append(cirq.Rz(rads=phi).on_each(*qubits))
+        emit.circuit.append(cirq.Rz(rads=phi).on_each(*qubits))
 
         return ()

src/bloqade/cirq_utils/emit/noise.py

@@ -34,7 +34,7 @@ def depolarize(
         p = frame.get(stmt.p)
         qubits = frame.get(stmt.qubits)
         cirfq_op = cirq.depolarize(p, n_qubits=1).on_each(qubits)
-        frame.circuit.append(cirfq_op)
+        interp.circuit.append(cirfq_op)
         return ()
 
     @impl(noise.stmts.Depolarize2)
@@ -46,7 +46,7 @@ def depolarize2(
         targets = frame.get(stmt.targets)
         cirq_qubits = [(ctrl, target) for ctrl, target in zip(controls, targets)]
         cirq_op = cirq.depolarize(p, n_qubits=2).on_each(cirq_qubits)
-        frame.circuit.append(cirq_op)
+        interp.circuit.append(cirq_op)
         return ()
 
     @impl(noise.stmts.SingleQubitPauliChannel)
@@ -62,7 +62,7 @@ def single_qubit_pauli_channel(
         qubits = frame.get(stmt.qubits)
 
         cirq_op = cirq.asymmetric_depolarize(px, py, pz).on_each(qubits)
-        frame.circuit.append(cirq_op)
+        interp.circuit.append(cirq_op)
 
         return ()
 
@@ -85,6 +85,6 @@ def two_qubit_pauli_channel(
         cirq_op = cirq.asymmetric_depolarize(
             error_probabilities=error_probabilities
         ).on_each(cirq_qubits)
-        frame.circuit.append(cirq_op)
+        interp.circuit.append(cirq_op)
 
         return ()

src/bloqade/cirq_utils/emit/qubit.py

@@ -23,13 +23,13 @@ def measure_qubit_list(
         self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Measure
     ):
         qbits = frame.get(stmt.qubits)
-        frame.circuit.append(cirq.measure(qbits))
+        emit.circuit.append(cirq.measure(qbits))
         return (emit.void,)
 
     @impl(qubit.Reset)
     def reset(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Reset):
         qubits = frame.get(stmt.qubits)
-        frame.circuit.append(
+        emit.circuit.append(
             cirq.ResetChannel().on_each(*qubits),
         )
         return ()

src/bloqade/cirq_utils/lowering.py

@@ -260,9 +260,7 @@ def run(
                 # NOTE: create a new register of appropriate size
                 n_qubits = len(self.qreg_index)
                 n = frame.push(py.Constant(n_qubits))
-                self.qreg = frame.push(
-                    func.Invoke((n.result,), callee=qalloc, kwargs=())
-                ).result
+                self.qreg = frame.push(func.Invoke((n.result,), callee=qalloc)).result
 
             self.visit(state, stmt)
 

src/bloqade/qasm2/emit/base.py

@@ -48,8 +48,7 @@ def initialize_frame(
     def run_method(
         self, method: ir.Method, args: tuple[ast.Node | None, ...]
     ) -> tuple[EmitQASM2Frame[StmtType], ast.Node | None]:
-        sym_name = method.sym_name if method.sym_name is not None else ""
-        return self.call(method, ast.Name(sym_name), *args)
+        return self.call(method, *args)
 
     def emit_block(self, frame: EmitQASM2Frame, block: ir.Block) -> ast.Node | None:
         for stmt in block.stmts: