bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/emit/gate.py

@@ -0,0 +1,102 @@
+from dataclasses import field, dataclass
+
+from kirin import ir, types, interp
+from kirin.dialects import py, func, ilist
+from kirin.ir.dialect import Dialect as Dialect
+
+from bloqade.types import QubitType
+from bloqade.qasm2.parse import ast
+
+from .base import EmitError, EmitQASM2Base, EmitQASM2Frame
+
+
+def _default_dialect_group():
+    from bloqade.qasm2.groups import gate
+
+    return gate
+
+
+@dataclass
+class EmitQASM2Gate(EmitQASM2Base[ast.UOp | ast.Barrier, ast.Gate]):
+    keys = ["emit.qasm2.gate"]
+    dialects: ir.DialectGroup = field(default_factory=_default_dialect_group)
+
+
+@ilist.dialect.register(key="emit.qasm2.gate")
+class Ilist(interp.MethodTable):
+
+    @interp.impl(ilist.New)
+    def emit_ilist(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: ilist.New):
+        return (ilist.IList(data=frame.get_values(stmt.values)),)
+
+
+@py.constant.dialect.register(key="emit.qasm2.gate")
+class Constant(interp.MethodTable):
+
+    @interp.impl(py.Constant)
+    def emit_constant(
+        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: py.Constant
+    ):
+        return (stmt.value,)
+
+
+@func.dialect.register(key="emit.qasm2.gate")
+class Func(interp.MethodTable):
+
+    @interp.impl(func.Call)
+    def emit_call(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Call):
+        raise EmitError("cannot emit dynamic call")
+
+    @interp.impl(func.Invoke)
+    def emit_invoke(
+        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Invoke
+    ):
+        ret = ()
+        if len(stmt.results) == 1 and stmt.results[0].type.is_subseteq(types.NoneType):
+            ret = (None,)
+        elif len(stmt.results) > 0:
+            raise EmitError(
+                "cannot emit invoke with results, this "
+                "is not compatible QASM2 gate routine"
+                " (consider pass qreg/creg by argument)"
+            )
+
+        cparams, qparams = [], []
+        for arg in stmt.inputs:
+            if arg.type.is_subseteq(QubitType):
+                qparams.append(frame.get(arg))
+            else:
+                cparams.append(frame.get(arg))
+
+        frame.body.append(
+            ast.Instruction(
+                name=ast.Name(stmt.callee.sym_name),
+                params=cparams,
+                qargs=qparams,
+            )
+        )
+        return ret
+
+    @interp.impl(func.Lambda)
+    @interp.impl(func.GetField)
+    def emit_err(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
+        raise EmitError(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 ()
+
+    @interp.impl(func.Function)
+    def emit_func(
+        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Function
+    ):
+        emit.run_ssacfg_region(frame, stmt.body)
+        cparams, qparams = [], []
+        for arg in stmt.args:
+            if arg.type.is_subseteq(QubitType):
+                qparams.append(frame.get(arg))
+            else:
+                cparams.append(frame.get(arg))
+        emit.output = ast.Gate(stmt.sym_name, cparams, qparams, frame.body)
+        return ()

bloqade/qasm2/emit/main.py

@@ -0,0 +1,106 @@
+from dataclasses import dataclass
+
+from kirin import ir, interp
+from kirin.dialects import cf, scf, func
+from kirin.ir.dialect import Dialect as Dialect
+
+from bloqade.qasm2.parse import ast
+
+from .base import EmitQASM2Base, EmitQASM2Frame
+
+
+@dataclass
+class EmitQASM2Main(EmitQASM2Base[ast.Statement, ast.MainProgram]):
+    keys = ["emit.qasm2.main", "emit.qasm2.gate"]
+    dialects: ir.DialectGroup
+
+
+@func.dialect.register(key="emit.qasm2.main")
+class Func(interp.MethodTable):
+
+    @interp.impl(func.Function)
+    def emit_func(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: func.Function
+    ):
+        from bloqade.qasm2.dialects import glob, noise, parallel
+
+        emit.run_ssacfg_region(frame, stmt.body)
+        if emit.dialects.data.intersection(
+            (parallel.dialect, glob.dialect, noise.dialect)
+        ):
+            header = ast.Kirin([dialect.name for dialect in emit.dialects])
+        else:
+            header = ast.OPENQASM(ast.Version(2, 0))
+
+        emit.output = ast.MainProgram(header=header, statements=frame.body)
+        return ()
+
+
+@cf.dialect.register(key="emit.qasm2.main")
+class Cf(interp.MethodTable):
+
+    @interp.impl(cf.Branch)
+    def emit_branch(self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: cf.Branch):
+        frame.worklist.append(
+            interp.Successor(stmt.successor, frame.get_values(stmt.arguments))
+        )
+        return ()
+
+    @interp.impl(cf.ConditionalBranch)
+    def emit_conditional_branch(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: cf.ConditionalBranch
+    ):
+        cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
+        body_frame = emit.new_frame(stmt)
+        body_frame.entries.update(frame.entries)
+        body_frame.set_values(
+            stmt.then_successor.args, frame.get_values(stmt.then_arguments)
+        )
+        emit.emit_block(body_frame, stmt.then_successor)
+        frame.body.append(
+            ast.IfStmt(
+                cond,
+                body=body_frame.body,  # type: ignore
+            )
+        )
+        frame.worklist.append(
+            interp.Successor(stmt.else_successor, frame.get_values(stmt.else_arguments))
+        )
+        return ()
+
+
+@scf.dialect.register(key="emit.qasm2.main")
+class Scf(interp.MethodTable):
+
+    @interp.impl(scf.Yield)
+    def emit_yield(self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: scf.Yield):
+        return frame.get_values(stmt.values)
+
+    @interp.impl(scf.IfElse)
+    def emit_if_else(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: scf.IfElse
+    ):
+        else_stmts = stmt.else_body.blocks[0].stmts
+        if not (
+            len(else_stmts) == 0
+            or len(else_stmts) == 1
+            and isinstance(else_stmts.at(0), scf.Yield)
+        ):
+            raise interp.InterpreterError(
+                "cannot lower if-else with non-empty else block"
+            )
+
+        cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
+        then_frame = emit.new_frame(stmt)
+        then_frame.entries.update(frame.entries)
+        emit.emit_block(then_frame, stmt.then_body.blocks[0])
+        frame.body.append(
+            ast.IfStmt(
+                cond,
+                body=then_frame.body,  # type: ignore
+            )
+        )
+        term = stmt.then_body.blocks[0].last_stmt
+        if isinstance(term, scf.Yield):
+            return then_frame.get_values(term.values)
+        return ()

bloqade/qasm2/emit/target.py

@@ -0,0 +1,165 @@
+import io
+
+from kirin import ir
+from rich.console import Console
+from kirin.analysis import CallGraph
+from kirin.dialects import ilist
+
+from bloqade.qasm2.parse import ast, pprint
+from bloqade.qasm2.passes.fold import QASM2Fold
+from bloqade.qasm2.passes.glob import GlobalToParallel
+from bloqade.qasm2.passes.py2qasm import Py2QASM
+from bloqade.qasm2.passes.parallel import ParallelToUOp
+
+from .gate import EmitQASM2Gate
+from .main import EmitQASM2Main
+
+
+class QASM2:
+    """QASM2 target for Bloqade kernels.
+
+    QASM2 target that accepts a Bloqade kernel and produces an AST that you can then obtain a string for printing or saving as a file.
+    """
+
+    def __init__(
+        self,
+        qelib1: bool = True,
+        allow_parallel: bool = False,
+        allow_global: bool = False,
+        custom_gate: bool = True,
+    ) -> None:
+        """Initialize the QASM2 target.
+
+        Args:
+            allow_parallel (bool):
+                Allow parallel gate in the resulting QASM2 AST. Defaults to `False`.
+                In the case its False, and the input kernel uses parallel gates, they will get rewrite into uop gates.
+
+            allow_global (bool):
+                Allow global gate in the resulting QASM2 AST. Defaults to `False`.
+                In the case its False, and the input kernel uses global gates, they will get rewrite into parallel gates.
+                If both `allow_parallel` and `allow_global` are False, the input kernel will be rewritten to use uop gates.
+
+            qelib1 (bool):
+                Include the `include "qelib1.inc"` line in the resulting QASM2 AST that's
+                submitted to qBraid. Defaults to `True`.
+            custom_gate (bool):
+                Include the custom gate definitions in the resulting QASM2 AST. Defaults to `True`. If `False`, all the qasm2.gate will be inlined.
+
+
+
+        """
+        from bloqade import qasm2
+
+        self.main_target = qasm2.main
+        self.gate_target = qasm2.gate
+
+        self.qelib1 = qelib1
+        self.custom_gate = custom_gate
+        self.allow_parallel = allow_parallel
+        self.allow_global = allow_global
+
+        if allow_parallel:
+            self.main_target = self.main_target.add(qasm2.dialects.parallel)
+            self.gate_target = self.gate_target.add(qasm2.dialects.parallel)
+
+        if allow_global:
+            self.main_target = self.main_target.add(qasm2.dialects.glob)
+            self.gate_target = self.gate_target.add(qasm2.dialects.glob)
+
+        if allow_global or allow_parallel:
+            self.main_target = self.main_target.add(ilist)
+            self.gate_target = self.gate_target.add(ilist)
+
+    def emit(self, entry: ir.Method) -> ast.MainProgram:
+        """Emit a QASM2 AST from the Bloqade kernel.
+
+        Args:
+            entry (ir.Method):
+                The Bloqade kernel to convert to the QASM2 AST
+
+        Returns:
+            ast.MainProgram:
+                A QASM2 AST object
+
+        """
+        assert len(entry.args) == 0, "entry method should not have arguments"
+
+        # make a cloned instance of kernel
+        entry = entry.similar()
+        QASM2Fold(entry.dialects, inline_gate_subroutine=not self.custom_gate).fixpoint(
+            entry
+        )
+
+        if not self.allow_global:
+            # rewrite global to parallel
+            GlobalToParallel(dialects=entry.dialects)(entry)
+
+        if not self.allow_parallel:
+            # rewrite parallel to uop
+            ParallelToUOp(dialects=entry.dialects)(entry)
+
+        Py2QASM(entry.dialects)(entry)
+        target_main = EmitQASM2Main(self.main_target)
+        target_main.run(
+            entry, tuple(ast.Name(name) for name in entry.arg_names[1:])
+        ).expect()
+
+        main_program = target_main.output
+        assert main_program is not None, f"failed to emit {entry.sym_name}"
+
+        extra = []
+        if self.qelib1:
+            extra.append(ast.Include("qelib1.inc"))
+
+        if self.custom_gate:
+            cg = CallGraph(entry)
+            target_gate = EmitQASM2Gate(self.gate_target)
+
+            for _, fn in cg.defs.items():
+                if fn is entry:
+                    continue
+
+                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_parallel:
+                    # rewrite parallel to uop
+                    ParallelToUOp(dialects=fn.dialects)(fn)
+
+                Py2QASM(fn.dialects)(fn)
+
+                target_gate.run(
+                    fn, tuple(ast.Name(name) for name in fn.arg_names[1:])
+                ).expect()
+                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
+
+    def emit_str(self, entry: ir.Method) -> str:
+        """Emit a QASM2 AST from the Bloqade kernel.
+
+        Args:
+            entry (ir.Method):
+                The Bloqade kernel to convert to the QASM2 AST
+
+        Returns:
+            str:
+                A string with the QASM2 representation of the kernel
+
+        """
+        console = Console(
+            file=io.StringIO(),
+            force_terminal=False,
+            force_interactive=False,
+            force_jupyter=False,
+            record=True,
+        )
+        pprint(self.emit(entry), console=console)
+        return console.export_text()

bloqade/qasm2/glob.py

@@ -0,0 +1,24 @@
+"""QASM2 extension for global gates."""
+
+from typing import Any
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from .types import QReg
+from .dialects import glob
+
+
+@wraps(glob.UGate)
+def u(
+    theta: float, phi: float, lam: float, registers: ilist.IList[QReg, Any] | list
+) -> None:
+    """Apply a U gate to all qubits in the input registers.
+
+    Args:
+        theta (float): The angle theta.
+        phi (float): The angle phi.
+        lam (float): The angle lam.
+        registers (IList[QReg] | list[QReg]): The registers to apply the gate to.
+
+    """

bloqade/qasm2/groups.py

@@ -0,0 +1,120 @@
+from kirin import ir, passes
+from kirin.prelude import structural_no_opt
+from kirin.dialects import scf, func, ilist, lowering
+
+from bloqade.qasm2.dialects import (
+    uop,
+    core,
+    expr,
+    glob,
+    noise,
+    inline,
+    indexing,
+    parallel,
+)
+from bloqade.qasm2.rewrite.desugar import IndexingDesugarPass
+
+
+@ir.dialect_group([uop, func, expr, lowering.func, lowering.call])
+def gate(self):
+    fold_pass = passes.Fold(self)
+    typeinfer_pass = passes.TypeInfer(self)
+
+    def run_pass(
+        method: ir.Method,
+        *,
+        fold: bool = True,
+    ):
+        method.verify()
+
+        if isinstance(method.code, func.Function):
+            new_code = expr.GateFunction(
+                sym_name=method.code.sym_name,
+                signature=method.code.signature,
+                body=method.code.body,
+            )
+            method.code = new_code
+        else:
+            raise ValueError(
+                "Gate Method code must be a Function, cannot be lambda/closure"
+            )
+
+        if fold:
+            fold_pass(method)
+
+        typeinfer_pass(method)
+        method.verify_type()
+
+    return run_pass
+
+
+@ir.dialect_group(
+    [
+        uop,
+        expr,
+        core,
+        scf,
+        indexing,
+        func,
+        lowering.func,
+        lowering.call,
+    ]
+)
+def main(self):
+    fold_pass = passes.Fold(self)
+    typeinfer_pass = passes.TypeInfer(self)
+
+    def run_pass(
+        method: ir.Method,
+        *,
+        fold: bool = True,
+    ):
+        method.verify()
+        # TODO make special Function rewrite
+
+        if fold:
+            fold_pass(method)
+
+        typeinfer_pass(method)
+        method.verify_type()
+
+    return run_pass
+
+
+@ir.dialect_group(
+    structural_no_opt.union(
+        [
+            inline,
+            uop,
+            glob,
+            noise,
+            parallel,
+            core,
+        ]
+    )
+)
+def extended(self):
+    fold_pass = passes.Fold(self)
+    typeinfer_pass = passes.TypeInfer(self)
+    ilist_desugar_pass = ilist.IListDesugar(self)
+    indexing_desugar_pass = IndexingDesugarPass(self)
+
+    def run_pass(
+        mt: ir.Method,
+        *,
+        fold: bool = True,
+        typeinfer: bool = True,
+    ):
+        mt.verify()
+        if fold:
+            fold_pass.fixpoint(mt)
+
+        if typeinfer:
+            typeinfer_pass(mt)
+        ilist_desugar_pass(mt)
+        indexing_desugar_pass(mt)
+        if typeinfer:
+            typeinfer_pass(mt)  # fix types after desugaring
+            mt.verify_type()
+
+    return run_pass

bloqade/qasm2/parallel.py

@@ -0,0 +1,48 @@
+"""QASM2 extension for parallel execution of gates."""
+
+from typing import Any
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from .types import Qubit
+from .dialects import parallel
+
+
+@wraps(parallel.CZ)
+def cz(
+    ctrls: ilist.IList[Qubit, Any] | list, qargs: ilist.IList[Qubit, Any] | list
+) -> None:
+    """Apply a controlled-Z gate to input qubits in parallel.
+
+    Args:
+        ctrls (IList[Qubit] | list[Qubit]): The control qubits.
+        qargs (IList[Qubit] | list[Qubit]): The target qubits.
+
+    """
+
+
+@wraps(parallel.UGate)
+def u(
+    qargs: ilist.IList[Qubit, Any] | list, theta: float, phi: float, lam: float
+) -> None:
+    """Apply a U gate to input qubits in parallel.
+
+    Args:
+        qargs (IList[Qubit] | list[Qubit]): The target qubits.
+        theta (float): The angle theta.
+        phi (float): The angle phi.
+        lam (float): The angle lam.
+
+    """
+
+
+@wraps(parallel.RZ)
+def rz(qargs: ilist.IList[Qubit, Any] | list, theta: float) -> None:
+    """Apply a RZ gate to input qubits in parallel.
+
+    Args:
+        qargs (IList[Qubit] | list[Qubit]): The target qubits.
+        theta (float): The angle theta.
+
+    """

bloqade/qasm2/parse/__init__.py

@@ -0,0 +1,37 @@
+import pathlib
+
+from rich.console import Console
+
+from . import ast as ast
+from .build import Build
+from .print import Printer as Printer
+from .parser import qasm2_parser as lark_parser
+from .visitor import Visitor as Visitor
+
+
+def loads(txt: str):
+    raw = lark_parser.parse(txt)
+    return Build().build_mainprogram(raw)
+
+
+def loadfile(file: str | pathlib.Path):
+    with open(file) as f:
+        return loads(f.read())
+
+
+def pprint(node: ast.Node, *, console: Console | None = None):
+    if console:
+        return Printer(console).visit(node)
+    else:
+        Printer().visit(node)
+
+
+def spprint(node: ast.Node, *, console: Console | None = None):
+    if console:
+        printer = Printer(console)
+    else:
+        printer = Printer()
+
+    with printer.string_io() as stream:
+        printer.visit(node)
+        return stream.getvalue()