Refactor QASM2 emission logic. Removed handling GateFunction, it is now handled in func.Function.

Author: zhenrongliew

src/bloqade/cirq_utils/emit/base.py

@@ -229,15 +229,9 @@ class __FuncEmit(MethodTable):
 
     @impl(func.Function)
     def emit_func(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Function):
-        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
+        result = emit.frame_eval(frame, stmt)
+        if isinstance(result, tuple):
+            frame.set_values(stmt.results, result)
 
         return (emit.circuit,)
 

src/bloqade/qasm2/dialects/expr/_emit.py

@@ -13,42 +13,6 @@
 @dialect.register(key="emit.qasm2.gate")
 class EmitExpr(interp.MethodTable):
 
-    @interp.impl(stmts.GateFunction)
-    def emit_func(
-        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.GateFunction
-    ):
-
-        args: list[ast.Node] = []
-        cparams, qparams = [], []
-        entry_args = stmt.body.blocks[0].args
-        user_args = entry_args[1:] if len(entry_args) > 0 else []
-
-        for arg in user_args:
-            assert arg.name is not None
-
-            args.append(ast.Name(id=arg.name))
-            if arg.type.is_subseteq(QubitType):
-                qparams.append(arg.name)
-            else:
-                cparams.append(arg.name)
-
-        frame.worklist.append(interp.Successor(stmt.body.blocks[0], *args))
-        if len(entry_args) > 0:
-            frame.set(entry_args[0], ast.Name(stmt.sym_name or "gate"))
-
-        while (succ := frame.worklist.pop()) is not None:
-            frame.set_values(succ.block.args[1:], succ.block_args)
-            block_header = emit.emit_block(frame, succ.block)
-            frame.block_ref[succ.block] = block_header
-        return (
-            ast.Gate(
-                name=stmt.sym_name,
-                cparams=cparams,
-                qparams=qparams,
-                body=frame.body,
-            ),
-        )
-
     @interp.impl(stmts.ConstInt)
     @interp.impl(stmts.ConstFloat)
     def emit_const_int(

src/bloqade/qasm2/emit/base.py

@@ -53,18 +53,17 @@ def run_method(
     def emit_block(self, frame: EmitQASM2Frame, block: ir.Block) -> ast.Node | None:
         for stmt in block.stmts:
             result = self.frame_eval(frame, stmt)
+            print(stmt.results, result)
             if isinstance(result, tuple):
-                if len(result) == 0:
-                    continue
-                keys = getattr(stmt, "_results", None) or getattr(stmt, "results", None)
-                if keys is None:
-                    continue
-                frame.set_values(keys, result)
+                frame.set_values(stmt.results, result)
         return None
 
     A = TypeVar("A")
     B = TypeVar("B")
 
+    def eval_fallback(self, frame: EmitQASM2Frame, node: ir.Statement):
+        return tuple(None for _ in range(len(node.results)))
+
     @overload
     def assert_node(self, typ: type[A], node: ast.Node | None) -> A: ...
 

src/bloqade/qasm2/emit/gate.py

@@ -1,6 +1,7 @@
 from dataclasses import field, dataclass
 
 from kirin import ir, types, interp
+from kirin.ir import method
 from kirin.dialects import py, func, ilist
 from kirin.ir.dialect import Dialect as Dialect
 from typing_extensions import Self
@@ -87,6 +88,5 @@ def emit_err(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         raise RuntimeError(f"illegal statement {stmt.name} for QASM2 gate routine")
 
     @interp.impl(func.Return)
-    @interp.impl(func.ConstantNone)
     def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         return ()

src/bloqade/qasm2/emit/main.py

@@ -6,6 +6,7 @@
 from kirin.ir.dialect import Dialect as Dialect
 from typing_extensions import Self
 
+from bloqade.types import QubitType
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.dialects.uop import SingleQubitGate, TwoQubitCtrlGate
 from bloqade.qasm2.dialects.expr import GateFunction
@@ -105,27 +106,40 @@ def emit_func(
                     with gate_emitter.new_frame(
                         callable_node, has_parent_access=False
                     ) as gate_frame:
-                        gate_result = gate_emitter.frame_eval(gate_frame, callable_node)
-                        gate_obj = None
-                        if isinstance(gate_result, tuple) and len(gate_result) > 0:
-                            maybe = gate_result[0]
-                            if isinstance(maybe, ast.Gate):
-                                gate_obj = maybe
-
-                        if gate_obj is None:
-                            name = emit.callables.get(
-                                callable_node
-                            ) or emit.callables.add(callable_node)
-                            prefix = getattr(emit.callables, "prefix", "") or ""
-                            emit_name = (
-                                name[len(prefix) :]
-                                if prefix and name.startswith(prefix)
-                                else name
-                            )
-                            gate_obj = ast.Gate(
-                                name=emit_name, cparams=[], qparams=[], body=[]
+                        args: list[ast.Node] = []
+                        cparams, qparams = [], []
+                        entry_args = callable_node.body.blocks[0].args
+                        user_args = entry_args[1:] if len(entry_args) > 0 else []
+
+                        for arg in user_args:
+                            assert arg.name is not None
+
+                            args.append(ast.Name(id=arg.name))
+                            if arg.type.is_subseteq(QubitType):
+                                qparams.append(arg.name)
+                            else:
+                                cparams.append(arg.name)
+
+                        # Map block arguments to AST names in the gate frame
+                        for arg in user_args:
+                            gate_frame.set(
+                                arg, ast.Name(id=getattr(arg, "name", "arg"))
                             )
 
+                        # Actually emit the gate body by interpreting all blocks
+                        for block in callable_node.body.blocks:
+                            gate_emitter.emit_block(gate_frame, block)
+
+                        name = emit.callables.get(callable_node) or emit.callables.add(
+                            callable_node
+                        )
+                        gate_obj = ast.Gate(
+                            name=name,
+                            cparams=cparams,
+                            qparams=qparams,
+                            body=gate_frame.body,
+                        )
+
                         gate_defs.append(gate_obj)
 
         if emit.dialects.data.intersection((parallel.dialect, glob.dialect)):

src/bloqade/qasm2/emit/target.py

@@ -124,10 +124,8 @@ def emit(self, entry: ir.Method) -> ast.MainProgram:
         extra = []
         if self.qelib1:
             extra.append(ast.Include("qelib1.inc"))
-
         if self.custom_gate:
             cg = CallGraph(entry)
-            target_gate = EmitQASM2Gate(self.gate_target).initialize()
 
             for _, fns in cg.defs.items():
                 if len(fns) != 1:
@@ -140,20 +138,16 @@ def emit(self, entry: ir.Method) -> ast.MainProgram:
                 fn = fn.similar()
                 QASM2Fold(fn.dialects).fixpoint(fn)
 
-                if not self.allow_global:
-                    # rewrite global to parallel
-                    GlobalToParallel(dialects=fn.dialects)(fn)
+                # if not self.allow_global:
+                #     # rewrite global to parallel
+                #     GlobalToParallel(dialects=fn.dialects)(fn)
 
-                if not self.allow_parallel:
-                    # rewrite parallel to uop
-                    ParallelToUOp(dialects=fn.dialects)(fn)
+                # if not self.allow_parallel:
+                #     # rewrite parallel to uop
+                #     ParallelToUOp(dialects=fn.dialects)(fn)
 
                 Py2QASM(fn.dialects)(fn)
 
-                target_gate.run(fn)
-                assert target_gate.output is not None, f"failed to emit {fn.sym_name}"
-                extra.append(target_gate.output)
-
         main_program.statements = extra + main_program.statements
         return main_program
 

test/qasm2/emit/test_qasm2.py

@@ -1,3 +1,4 @@
+import io
 from pathlib import Path
 from contextlib import redirect_stdout
 
@@ -30,14 +31,9 @@ def main():
     ast = target.emit(main)
     filename = "t_qasm2.qasm"
     # Read the filename and store in target
-    with open(Path(__file__).parent.resolve() / filename, "r") as io:
-        target = io.read()
-    # Output into filename in the current directory
-    out_path = Path(__file__).parent.resolve() / filename
-    with open(out_path, "w") as f:
-        with redirect_stdout(f):
-            qasm2.parse.pprint(ast)
-    # Read the filename again and store in generated
-    with open(Path(__file__).parent.resolve() / filename, "r") as io:
-        generated = io.read()
-    assert generated.strip() == target.strip()
+    with open(Path(__file__).parent.resolve() / filename, "r") as txt:
+        target = txt.read()
+    buf = io.StringIO()
+    with redirect_stdout(buf):
+        qasm2.parse.pprint(ast)
+    assert buf.getvalue().strip() == target.strip()