Add back returns kwarg to qasm2.loads / qasm2.loadfile (#266)

Author: david-pl

src/bloqade/qasm2/_qasm_loading.py

@@ -3,7 +3,7 @@
 import pathlib
 from typing import Any
 
-from kirin import ir, types
+from kirin import ir, lowering
 from kirin.dialects import func
 
 from . import parse
@@ -16,6 +16,7 @@ def loads(
     *,
     kernel_name: str = "main",
     dialects: ir.DialectGroup | None = None,
+    returns: str | None = None,
     globals: dict[str, Any] | None = None,
     file: str | None = None,
     lineno_offset: int = 0,
@@ -54,7 +55,7 @@ def loads(
     # TODO: add source info
     stmt = parse.loads(qasm)
     qasm2_lowering = QASM2(dialects or main)
-    body = qasm2_lowering.run(
+    frame = qasm2_lowering.get_frame(
         stmt,
         source=qasm,
         file=file,
@@ -63,13 +64,21 @@ def loads(
         col_offset=col_offset,
         compactify=compactify,
     )
-    return_value = func.ConstantNone()
-    body.blocks[0].stmts.append(return_value)
-    body.blocks[0].stmts.append(func.Return(value_or_stmt=return_value))
 
+    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))
+
+    body = frame.curr_region
     code = func.Function(
         sym_name=kernel_name,
-        signature=func.Signature((), types.NoneType),
+        signature=func.Signature((), return_node.value.type),
         body=body,
     )
 
@@ -88,6 +97,7 @@ def loadfile(
     *,
     kernel_name: str = "main",
     dialects: ir.DialectGroup | None = None,
+    returns: str | None = None,
     globals: dict[str, Any] | None = None,
     file: str | None = None,
     lineno_offset: int = 0,
@@ -132,6 +142,7 @@ def loadfile(
         source,
         kernel_name=kernel_name,
         dialects=dialects,
+        returns=returns,
         globals=globals,
         file=file,
         lineno_offset=lineno_offset,

src/bloqade/qasm2/parse/lowering.py

@@ -28,6 +28,28 @@ def run(
         col_offset: int = 0,
         compactify: bool = True,
     ) -> ir.Region:
+
+        frame = self.get_frame(
+            stmt,
+            source=source,
+            globals=globals,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+
+        return frame.curr_region
+
+    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,
+    ) -> lowering.Frame:
         # TODO: add source info
         state = lowering.State(
             self,
@@ -41,13 +63,13 @@ def run(
             finalize_next=False,
         ) as frame:
             self.visit(state, stmt)
-            region = frame.curr_region
 
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
+            if compactify:
+                from kirin.rewrite import Walk, CFGCompactify
+
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
 
-            Walk(CFGCompactify()).rewrite(region)
-        return region
+            return frame
 
     def visit(self, state: lowering.State[ast.Node], node: ast.Node) -> lowering.Result:
         name = node.__class__.__name__

test/pyqrack/test_target.py

@@ -1,11 +1,12 @@
 import math
+import textwrap
 
 import pytest
 from kirin import ir
 from kirin.dialects import ilist
 
 from bloqade import qasm2
-from bloqade.pyqrack import PyQrack, PyQrackQubit, StackMemorySimulator, reg
+from bloqade.pyqrack import PyQrack, CRegister, PyQrackQubit, StackMemorySimulator, reg
 
 
 def test_target():
@@ -171,3 +172,45 @@ def parallel():
     result = target.run(parallel)
 
     assert result == [reg.Measurement.One] * 4
+
+
+def test_loads_without_return():
+    qasm2_str = textwrap.dedent(
+        """
+    OPENQASM 2.0;
+
+    qreg q[1];
+    x q[0];
+    """
+    )
+
+    main = qasm2.loads(qasm2_str)
+
+    sim = StackMemorySimulator(min_qubits=1)
+
+    result = sim.run(main)
+    assert result is None
+
+    ket = sim.state_vector(main)
+    assert ket[0] == 0
+
+
+def test_loads_with_return():
+    qasm2_str = textwrap.dedent(
+        """
+    OPENQASM 2.0;
+
+    qreg q[1];
+    creg c[1];
+    x q[0];
+    measure q -> c;
+    """
+    )
+
+    main = qasm2.loads(qasm2_str, returns="c")
+
+    sim = StackMemorySimulator(min_qubits=1)
+    result = sim.run(main)
+
+    assert isinstance(result, CRegister)
+    assert result[0] == 1