Code Generation for QASM2 (#555)

Refactor QASM2's code generation to use the new Emit APIs.

---------

Co-authored-by: David Plankensteiner <[email protected]>

Author: zhenrongliew

src/bloqade/analysis/fidelity/analysis.py

@@ -83,12 +83,15 @@ def run_analysis(
         self, method: ir.Method, args: tuple | None = None, *, no_raise: bool = True
     ) -> tuple[ForwardFrame, Any]:
         self._run_address_analysis(method, no_raise=no_raise)
-        return super().run_analysis(method, args, no_raise=no_raise)
+        return super().run(method)
 
     def _run_address_analysis(self, method: ir.Method, no_raise: bool):
         addr_analysis = AddressAnalysis(self.dialects)
-        addr_frame, _ = addr_analysis.run_analysis(method=method, no_raise=no_raise)
+        addr_frame, _ = addr_analysis.run(method=method)
         self.addr_frame = addr_frame
 
         # NOTE: make sure we have as many probabilities as we have addresses
         self.atom_survival_probability = [1.0] * addr_analysis.qubit_count
+
+    def method_self(self, method: ir.Method) -> EmptyLattice:
+        return self.lattice.bottom()

src/bloqade/analysis/measure_id/analysis.py

@@ -53,3 +53,6 @@ def get_const_value(
                 return hint.data
 
         return None
+
+    def method_self(self, method: ir.Method) -> MeasureId:
+        return self.lattice.bottom()

src/bloqade/cirq_utils/emit/base.py

@@ -4,9 +4,9 @@
 
 import cirq
 from kirin import ir, types, interp
-from kirin.emit import EmitABC, EmitError, EmitFrame
+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
@@ -102,7 +102,7 @@ def main():
         and isinstance(mt.code, func.Function)
         and not mt.code.signature.output.is_subseteq(types.NoneType)
     ):
-        raise EmitError(
+        raise interp.exceptions.InterpreterError(
             "The method you are trying to convert to a circuit has a return value, but returning from a circuit is not supported."
             " Set `ignore_returns = True` in order to simply ignore the return values and emit a circuit."
         )
@@ -116,12 +116,14 @@ def main():
 
     symbol_op_trait = mt.code.get_trait(ir.SymbolOpInterface)
     if (symbol_op_trait := mt.code.get_trait(ir.SymbolOpInterface)) is None:
-        raise EmitError("The method is not a symbol, cannot emit circuit!")
+        raise interp.exceptions.InterpreterError(
+            "The method is not a symbol, cannot emit circuit!"
+        )
 
     sym_name = symbol_op_trait.get_sym_name(mt.code).unwrap()
 
     if (signature_trait := mt.code.get_trait(ir.HasSignature)) is None:
-        raise EmitError(
+        raise interp.exceptions.InterpreterError(
             f"The method {sym_name} does not have a signature, cannot emit circuit!"
         )
 
@@ -135,7 +137,7 @@ def main():
 
     assert first_stmt is not None, "Method has no statements!"
     if len(args_ssa) - 1 != len(args):
-        raise EmitError(
+        raise interp.exceptions.InterpreterError(
             f"The method {sym_name} takes {len(args_ssa) - 1} arguments, but you passed in {len(args)} via the `args` keyword!"
         )
 
@@ -147,17 +149,22 @@ 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(
+        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():
@@ -166,23 +173,24 @@ def _default_kernel():
 
 @dataclass
 class EmitCirq(EmitABC[EmitCirqFrame, cirq.Circuit]):
-    keys = ["emit.cirq", "main"]
+    keys = ("emit.cirq", "emit.main")
     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,
@@ -196,7 +204,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
@@ -206,33 +214,43 @@ 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
+
+    def eval_fallback(self, frame: EmitCirqFrame, node: ir.Statement) -> tuple:
+        return tuple(None for _ in range(len(node.results)))
 
 
 @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):
-        raise EmitError(
+        raise interp.exceptions.InterpreterError(
             "Function invokes should need to be inlined! "
             "If you called the emit_circuit method, that should have happened, please report this issue."
         )
 
-    @impl(func.Return)
-    def return_(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Return):
-        # NOTE: should only be hit if ignore_returns == True
-        return ()
-
 
 @py.indexing.dialect.register(key="emit.cirq")
 class __Concrete(interp.MethodTable):
@@ -241,3 +259,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/native/upstream/__init__.py

@@ -1,5 +1,4 @@
 from .squin2native import (
     GateRule as GateRule,
     SquinToNative as SquinToNative,
-    SquinToNativePass as SquinToNativePass,
 )

src/bloqade/pyqrack/device.py

@@ -353,7 +353,7 @@ def task(
             kwargs = {}
 
         address_analysis = AddressAnalysis(dialects=kernel.dialects)
-        frame, _ = address_analysis.run_analysis(kernel)
+        frame, _ = address_analysis.run(kernel)
         if self.min_qubits == 0 and any(
             isinstance(a, (UnknownQubit, UnknownReg)) for a in frame.entries.values()
         ):

src/bloqade/pyqrack/target.py

@@ -51,7 +51,7 @@ def _get_interp(self, mt: ir.Method[Params, RetType]):
             return PyQrackInterpreter(mt.dialects, memory=DynamicMemory(options))
         else:
             address_analysis = AddressAnalysis(mt.dialects)
-            frame, _ = address_analysis.run_analysis(mt)
+            frame, _ = address_analysis.run(mt)
             if self.min_qubits == 0 and any(
                 isinstance(a, UnknownQubit) for a in frame.entries.values()
             ):