Merge branch 'main' into david/pyright-athon

Author: Roger-luo

src/bloqade/qasm2/parse/lowering.py

@@ -1,3 +1,5 @@
+import os
+import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -17,6 +19,118 @@ 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(
+        self,
+        source: str,
+        kernel_name: str,
+        *,
+        returns: 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)
+
+        state = lowering.State(
+            self,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+        with state.frame(
+            [stmt],
+            globals=globals,
+        ) 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()
+
+                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,
+        )
+
+    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(
         self,
         stmt: ast.Node,

test/qasm2/test_lowering.py

@@ -1,5 +1,9 @@
+import pathlib
+import tempfile
 import textwrap
 
+from kirin.dialects import func
+
 from bloqade import qasm2
 from bloqade.qasm2.parse.lowering import QASM2
 
@@ -17,6 +21,24 @@
 rx(pi/2) q[0];
 """
 )
-ast = qasm2.parse.loads(lines)
-code = QASM2(qasm2.main).run(ast)
-code.print()
+
+
+def test_run_lowering():
+    ast = qasm2.parse.loads(lines)
+    code = QASM2(qasm2.main).run(ast)
+    code.print()
+
+
+def test_loadfile():
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        with open(f"{tmp_dir}/test.qasm", "w") as f:
+            f.write(lines)
+
+        file = pathlib.Path(f"{tmp_dir}/test.qasm")
+        kernel = QASM2(qasm2.main).loadfile(file, returns="c")
+
+        assert isinstance(
+            (ret := kernel.callable_region.blocks[0].last_stmt), func.Return
+        )
+        assert ret.value.type.is_equal(qasm2.types.CRegType)

test/qasm2/test_native.py

@@ -2,7 +2,6 @@
 import textwrap
 
 import cirq
-import cirq.contrib
 import cirq.testing
 import cirq.contrib.qasm_import as qasm_import
 import cirq.circuits.qasm_output as qasm_output
@@ -86,6 +85,18 @@ def generator(n_tests: int):
     """
     )
 
+    yield textwrap.dedent(
+        """
+    OPENQASM 2.0;
+    include "qelib1.inc";
+
+    qreg q[2];
+
+    cu3(0.0, 0.6, 3.141591) q[0],q[1];
+
+    """
+    )
+
     rgen = np.random.RandomState(128)
     for num in range(n_tests):
         # Generate a new instance:
@@ -117,3 +128,28 @@ def kernel():
     cirq.testing.assert_allclose_up_to_global_phase(
         cirq.unitary(old_circuit), cirq.unitary(cirq_circuit), atol=1e-8
     )
+
+
+def test_cu3_rewrite():
+    prog = textwrap.dedent(
+        """
+        OPENQASM 2.0;
+        include "qelib1.inc";
+
+        qreg q[2];
+
+        cu3(0.0, 0.6, 3.141591) q[0],q[1];
+
+        """
+    )
+
+    @qasm2.main.add(qasm2.dialects.inline)
+    def kernel():
+        qasm2.inline(prog)
+
+    walk.Walk(RydbergGateSetRewriteRule(kernel.dialects)).rewrite(kernel.code)
+
+    new_qasm2 = qasm2.emit.QASM2().emit_str(kernel)
+
+    # simple-stupid test to see if the rewrite injected a bunch of new lines
+    assert new_qasm2.count("\n") > prog.count("\n")