bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/parse/print.py

@@ -0,0 +1,293 @@
+from dataclasses import field, dataclass
+
+from kirin import print
+
+from .ast import (
+    OPENQASM,
+    Pi,
+    Bit,
+    Cmp,
+    Call,
+    CReg,
+    Gate,
+    Name,
+    QReg,
+    BinOp,
+    Kirin,
+    Reset,
+    UGate,
+    CXGate,
+    IfStmt,
+    Number,
+    Opaque,
+    Barrier,
+    Comment,
+    Include,
+    Measure,
+    UnaryOp,
+    GlobUGate,
+    ParaCZGate,
+    ParaRZGate,
+    ParaU3Gate,
+    Instruction,
+    MainProgram,
+    NoisePAULI1,
+    ParallelQArgs,
+)
+from .visitor import Visitor
+
+
+@dataclass
+class ColorScheme:
+    comment: str = "bright_black"
+    keyword: str = "red"
+    symbol: str = "cyan"
+    string: str = "yellow"
+    number: str = "green"
+    irrational: str = "magenta"
+
+
+@dataclass
+class PrintState:
+    indent: int = 0
+    result_width: int = 0
+    rich_style: str | None = None
+    rich_highlight: bool | None = False
+    indent_marks: list[int] = field(default_factory=list)
+    messages: list[str] = field(default_factory=list)
+
+
+class Printer(print.Printer, Visitor[None]):
+
+    def visit_MainProgram(self, node: MainProgram) -> None:
+        self.print_indent()
+        self.visit(node.header)
+        self.print_newline()
+        for stmt in node.statements:
+            self.visit(stmt)
+            self.print_newline()
+
+    def visit_OPENQASM(self, node: OPENQASM) -> None:
+        self.plain_print(
+            f"OPENQASM {node.version.major}.{node.version.minor}",
+            style="comment",
+        )
+        self.plain_print(";")
+
+    def visit_Kirin(self, node: Kirin) -> None:
+        self.plain_print(
+            "KIRIN " + "{" + ",".join(sorted(node.dialects)) + "}", style="comment"
+        )
+        self.plain_print(";")
+
+    def visit_Include(self, node: Include) -> None:
+        self.plain_print("include", style="keyword")
+        self.plain_print(" ")
+        self.plain_print('"', node.filename, '"', style="string")
+        self.plain_print(";")
+
+    def visit_Barrier(self, node: Barrier) -> None:
+        self.print_indent()
+        self.plain_print("barrier", style="keyword")
+        self.plain_print(" ")
+        self.print_seq(node.qargs, emit=self.visit)
+        self.plain_print(";")
+
+    def visit_Instruction(self, node: Instruction) -> None:
+        self.visit_Name(node.name)
+        self.plain_print(" ")
+        if node.params:
+            self.print_seq(
+                node.params, delim=", ", prefix="(", suffix=") ", emit=self.visit
+            )
+        self.print_seq(node.qargs, emit=self.visit)
+        self.plain_print(";")
+
+    def visit_Comment(self, node: Comment) -> None:
+        self.plain_print("// ", node.text, style="comment")
+
+    def visit_CReg(self, node: CReg) -> None:
+        self.plain_print("creg", style="keyword")
+        self.plain_print(f" {node.name}[{node.size}]")
+        self.plain_print(";")
+
+    def visit_QReg(self, node: QReg) -> None:
+        self.plain_print("qreg", style="keyword")
+        self.plain_print(f" {node.name}[{node.size}]")
+        self.plain_print(";")
+
+    def visit_CXGate(self, node: CXGate) -> None:
+        self.plain_print("CX", style="keyword")
+        self.plain_print(" ")
+        self.visit(node.ctrl)
+        self.plain_print(", ")
+        self.visit(node.qarg)
+        self.plain_print(";")
+
+    def visit_UGate(self, node: UGate) -> None:
+        self.plain_print("U", style="keyword")
+        self.plain_print("(")
+        self.visit(node.theta)
+        self.plain_print(", ")
+        self.visit(node.phi)
+        self.plain_print(", ")
+        self.visit(node.lam)
+        self.plain_print(") ")
+        self.visit(node.qarg)
+        self.plain_print(";")
+
+    def visit_Measure(self, node: Measure) -> None:
+        self.plain_print("measure", style="keyword")
+        self.plain_print(" ")
+        self.visit(node.qarg)
+        self.plain_print(" -> ")
+        self.visit(node.carg)
+        self.plain_print(";")
+
+    def visit_Reset(self, node: Reset) -> None:
+        self.plain_print("reset ")
+        self.visit(node.qarg)
+        self.plain_print(";")
+
+    def visit_Opaque(self, node: Opaque) -> None:
+        self.plain_print("opaque ", style="keyword")
+        if node.cparams:
+            self.print_seq(
+                node.cparams, delim=", ", prefix="(", suffix=")", emit=self.visit
+            )
+
+        if node.qparams:
+            self.plain_print(" ")
+            self.print_seq(node.qparams, delim=", ", emit=self.visit)
+        self.plain_print(";")
+
+    def visit_Gate(self, node: Gate) -> None:
+        self.plain_print("gate ", style="keyword")
+        self.plain_print(node.name, style="symbol")
+        if node.cparams:
+            self.print_seq(
+                node.cparams, delim=", ", prefix="(", suffix=")", emit=self.plain_print
+            )
+
+        if node.qparams:
+            self.plain_print(" ")
+            self.print_seq(node.qparams, delim=", ", emit=self.plain_print)
+
+        self.plain_print(" {")
+        with self.indent():
+            self.print_newline()
+            for idx, stmt in enumerate(node.body):
+                self.visit(stmt)
+                if idx < len(node.body) - 1:
+                    self.print_newline()
+        self.print_newline()
+        self.plain_print("}")
+
+    def visit_IfStmt(self, node: IfStmt) -> None:
+        self.plain_print("if", style="keyword")
+        self.visit(node.cond)
+        if len(node.body) == 1:  # inline if
+            self.visit(node.body[0])
+        else:
+            self.plain_print("{")
+            with self.indent():
+                self.print_newline()
+                for idx, stmt in enumerate(node.body):
+                    self.visit(stmt)
+                    if idx < len(node.body) - 1:
+                        self.print_newline()
+            self.print_newline()
+            self.plain_print("}")
+
+    def visit_Cmp(self, node: Cmp) -> None:
+        self.plain_print(" (")
+        self.visit(node.lhs)
+        self.plain_print(" == ", style="keyword")
+        self.visit(node.rhs)
+        self.plain_print(") ")
+
+    def visit_Call(self, node: Call) -> None:
+        self.plain_print(node.name)
+        self.print_seq(node.args, delim=", ", prefix="(", suffix=")", emit=self.visit)
+
+    def visit_BinOp(self, node: BinOp) -> None:
+        self.plain_print("(")
+        self.visit(node.lhs)
+        self.plain_print(f" {node.op} ", style="keyword")
+        self.visit(node.rhs)
+        self.plain_print(")")
+
+    def visit_UnaryOp(self, node: UnaryOp) -> None:
+        self.plain_print(f"{node.op}", style="keyword")
+        self.visit(node.operand)
+
+    def visit_Bit(self, node: Bit) -> None:
+        self.visit_Name(node.name)
+        if node.addr is not None:
+            self.plain_print("[")
+            self.plain_print(node.addr, style="number")
+            self.plain_print("]")
+
+    def visit_Number(self, node: Number) -> None:
+        self.plain_print(node.value)
+
+    def visit_Pi(self, node: Pi) -> None:
+        self.plain_print("pi", style="number")
+
+    def visit_Name(self, node: Name) -> None:
+        return self.plain_print(node.id, style="symbol")
+
+    def visit_ParallelQArgs(self, node: ParallelQArgs) -> None:
+        self.plain_print("{")
+        with self.indent():
+            for idx, qargs in enumerate(node.qargs):
+                self.print_newline()
+                self.print_seq(qargs, emit=self.visit)
+                self.plain_print(";")
+        self.print_newline()
+        self.plain_print("}")
+
+    def visit_ParaU3Gate(self, node: ParaU3Gate) -> None:
+        self.plain_print("parallel.U", style="keyword")
+        self.plain_print("(")
+        self.visit(node.theta)
+        self.plain_print(", ")
+        self.visit(node.phi)
+        self.plain_print(", ")
+        self.visit(node.lam)
+        self.plain_print(") ")
+        self.visit_ParallelQArgs(node.qargs)
+
+    def visit_ParaCZGate(self, node: ParaCZGate) -> None:
+        self.plain_print("parallel.CZ ", style="keyword")
+        self.visit_ParallelQArgs(node.qargs)
+
+    def visit_ParaRZGate(self, node: ParaRZGate) -> None:
+        self.plain_print("parallel.RZ", style="keyword")
+        self.plain_print("(")
+        self.visit(node.theta)
+        self.plain_print(") ")
+        self.visit_ParallelQArgs(node.qargs)
+
+    def visit_GlobUGate(self, node: GlobUGate) -> None:
+        self.plain_print("glob.U", style="keyword")
+        self.plain_print("(")
+        self.visit(node.theta)
+        self.plain_print(", ")
+        self.visit(node.phi)
+        self.plain_print(", ")
+        self.visit(node.lam)
+        self.plain_print(") ")
+        self.print_seq(node.registers, prefix="{", suffix="}", emit=self.visit)
+
+    def visit_NoisePAULI1(self, node: NoisePAULI1) -> None:
+        self.plain_print("noise.PAULI1", style="keyword")
+        self.plain_print("(")
+        self.visit(node.px)
+        self.plain_print(", ")
+        self.visit(node.py)
+        self.plain_print(", ")
+        self.visit(node.pz)
+        self.plain_print(") ")
+        self.visit(node.qarg)
+        self.plain_print(";")

bloqade/qasm2/parse/qasm2.lark

@@ -0,0 +1,75 @@
+mainprogram: header ";" [statement*]
+?header: openqasm | kirin
+openqasm: "OPENQASM" version
+kirin: "KIRIN" "{" dialect ("," dialect)* "}"
+dialect: IDENTIFIER ("." IDENTIFIER)*
+version: INT "." INT
+?statement: regdecl | gate | opaque | qop | ifstmt | barrier | include
+?regdecl: qreg | creg
+// stmts
+include: "include" STRING ";"
+ifstmt: "if" "(" expr "==" expr ")" ifbody
+ifbody: qop | "{" qop* "}" // allow multiple qops
+opaque: "opaque" IDENTIFIER ["(" [params] ")"] qubits ";"
+barrier: "barrier" qubits ";"
+qreg: "qreg" IDENTIFIER "[" INT "]" ";"
+creg: "creg" IDENTIFIER "[" INT "]" ";"
+gate: "gate" IDENTIFIER ["(" cparams? ")"] qparams "{" (uop | barrier)* "}"
+cparams: IDENTIFIER ("," IDENTIFIER)*
+qparams: IDENTIFIER ("," IDENTIFIER)*
+
+// quantum ops
+?qop: uop | measure | extension | reset
+reset: "reset" bit ";"
+measure: "measure" bit "->" bit ";"
+
+?uop: inst | ugate | cx_gate
+inst: IDENTIFIER ["(" [params] ")"] qubits ";"
+ugate: "U" "(" expr "," expr "," expr ")" bit ";"
+cx_gate: "CX" bit "," bit ";"
+params: expr ("," expr)*
+
+?extension: parallel | glob | noise
+?glob: glob_u_gate
+glob_u_gate: "glob" "." "U" "(" expr "," expr "," expr ")" global_body
+global_body: "{" IDENTIFIER ("," IDENTIFIER)* "}"
+?noise: noise_pauli1
+noise_pauli1: "noise" "." "PAULI1" "(" expr "," expr "," expr ")" bit ";"
+?parallel: para_u_gate | para_rz_gate | para_cz_gate
+para_u_gate: "parallel" "." "U" "(" expr "," expr "," expr ")" parallel_body
+para_rz_gate: "parallel" "." "RZ" "(" expr ")" parallel_body
+para_cz_gate: "parallel" "." "CZ" parallel_body
+parallel_body: "{" task_args* "}"
+task_args: bit ("," bit)* ";"
+
+?bit: IDENTIFIER ("[" INT "]")?
+qubits: bit ("," bit)*
+
+?expr: term ((PLUS|MINUS) term)*
+?term: factor ((TIMES|DIVIDE) factor)*
+?factor: (PLUS|MINUS) factor | power
+?power: molecule (POW factor)?
+?molecule: molecule "(" [arglist] ")" -> call
+         | atom
+?atom: "(" expr ")" | IDENTIFIER | NUMBER | PI
+arglist: argument ("," argument)*
+?argument: expr
+
+VERSION: "2.0"
+PI: "pi"
+PLUS: "+"
+MINUS: "-"
+TIMES: "*"
+DIVIDE: "/"
+POW: "^"
+COMMENT: "//" /[^\n]*/ NEWLINE
+NEWLINE: "\n"
+
+%import common.INT
+%import common.FLOAT
+%import common.CNAME -> IDENTIFIER
+%import common.NUMBER
+%import common.ESCAPED_STRING -> STRING
+%import common.WS
+%ignore WS
+%ignore COMMENT

bloqade/qasm2/parse/visitor.py

@@ -0,0 +1,16 @@
+from typing import Generic, TypeVar
+
+from . import ast
+
+T = TypeVar("T")
+
+
+class Visitor(Generic[T]):
+
+    def visit(self, node: ast.Node) -> T:
+        method_name = f"visit_{node.__class__.__name__}"
+        visitor = getattr(self, method_name, self.generic_visit)
+        return visitor(node)
+
+    def generic_visit(self, node: ast.Node) -> T:
+        raise NotImplementedError(f"No visit_{node.__class__.__name__} method")

bloqade/qasm2/parse/visitor.pyi

@@ -0,0 +1,39 @@
+from typing import Generic, TypeVar
+
+from . import ast
+
+T = TypeVar("T")
+
+class Visitor(Generic[T]):
+    def visit(self, node: ast.Node) -> T: ...
+    def generic_visit(self, node: ast.Node) -> T: ...
+    def visit_MainProgram(self, node: ast.MainProgram) -> T: ...
+    def visit_OPENQASM(self, node: ast.OPENQASM) -> T: ...
+    def visit_Kirin(self, node: ast.Kirin) -> T: ...
+    def visit_QReg(self, node: ast.QReg) -> T: ...
+    def visit_CReg(self, node: ast.CReg) -> T: ...
+    def visit_Gate(self, node: ast.Gate) -> T: ...
+    def visit_Opaque(self, node: ast.Opaque) -> T: ...
+    def visit_IfStmt(self, node: ast.IfStmt) -> T: ...
+    def visit_Cmp(self, node: ast.Cmp) -> T: ...
+    def visit_Barrier(self, node: ast.Barrier) -> T: ...
+    def visit_Include(self, node: ast.Include) -> T: ...
+    def visit_Measure(self, node: ast.Measure) -> T: ...
+    def visit_Reset(self, node: ast.Reset) -> T: ...
+    def visit_Instruction(self, node: ast.Instruction) -> T: ...
+    def visit_UGate(self, node: ast.UGate) -> T: ...
+    def visit_CXGate(self, node: ast.CXGate) -> T: ...
+    def visit_Bit(self, node: ast.Bit) -> T: ...
+    def visit_BinOp(self, node: ast.BinOp) -> T: ...
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> T: ...
+    def visit_Call(self, node: ast.Call) -> T: ...
+    def visit_Number(self, node: ast.Number) -> T: ...
+    def visit_Pi(self, node: ast.Pi) -> T: ...
+    def visit_Name(self, node: ast.Name) -> T: ...
+    # extensions
+    def visit_ParaU3Gate(self, node: ast.ParaU3Gate) -> T: ...
+    def visit_ParaCZGate(self, node: ast.ParaCZGate) -> T: ...
+    def visit_ParaRZGate(self, node: ast.ParaRZGate) -> T: ...
+    def visit_ParallelQArgs(self, node: ast.ParallelQArgs) -> T: ...
+    def visit_GlobUGate(self, node: ast.GlobUGate) -> T: ...
+    def visit_NoisePAULI1(self, node: ast.NoisePAULI1) -> T: ...

bloqade/qasm2/passes/__init__.py

@@ -0,0 +1,5 @@
+from .fold import QASM2Fold as QASM2Fold
+from .noise import NoisePass as NoisePass
+from .py2qasm import Py2QASM as Py2QASM
+from .qasm2py import QASM2Py as QASM2Py
+from .parallel import UOpToParallel as UOpToParallel

bloqade/qasm2/passes/fold.py

@@ -0,0 +1,94 @@
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass, TypeInfer
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Inline,
+    Fixpoint,
+    WrapConst,
+    Call2Invoke,
+    ConstantFold,
+    CFGCompactify,
+    InlineGetItem,
+    InlineGetField,
+    DeadCodeElimination,
+    CommonSubexpressionElimination,
+)
+from kirin.analysis import const
+from kirin.dialects import scf, ilist
+from kirin.ir.method import Method
+from kirin.rewrite.abc import RewriteResult
+
+from bloqade.qasm2.dialects import expr
+
+
+@dataclass
+class QASM2Fold(Pass):
+    """Fold pass for qasm2.extended"""
+
+    constprop: const.Propagate = field(init=False)
+    inline_gate_subroutine: bool = True
+
+    def __post_init__(self):
+        self.constprop = const.Propagate(self.dialects)
+        self.typeinfer = TypeInfer(self.dialects)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        frame, _ = self.constprop.run_analysis(mt)
+        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
+        rule = Chain(
+            ConstantFold(),
+            Call2Invoke(),
+            InlineGetField(),
+            InlineGetItem(),
+            DeadCodeElimination(),
+            CommonSubexpressionElimination(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+        result = (
+            Walk(
+                Chain(
+                    scf.unroll.PickIfElse(),
+                    scf.unroll.ForLoop(),
+                    scf.trim.UnusedYield(),
+                )
+            )
+            .rewrite(mt.code)
+            .join(result)
+        )
+
+        # run typeinfer again after unroll etc. because we now insert
+        # a lot of new nodes, which might have more precise types
+        self.typeinfer.unsafe_run(mt)
+        result = (
+            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
+            .rewrite(mt.code)
+            .join(result)
+        )
+
+        def inline_simple(node: ir.Statement):
+            if isinstance(node, expr.GateFunction):
+                return self.inline_gate_subroutine
+
+            if not isinstance(node.parent_stmt, (scf.For, scf.IfElse)):
+                return True  # always inline calls outside of loops and if-else
+
+            # inside loops and if-else, only inline simple functions, i.e. functions with a single block
+            if (trait := node.get_trait(ir.CallableStmtInterface)) is None:
+                return False  # not a callable, don't inline to be safe
+            region = trait.get_callable_region(node)
+            return len(region.blocks) == 1
+
+        result = (
+            Walk(
+                Inline(inline_simple),
+            )
+            .rewrite(mt.code)
+            .join(result)
+        )
+        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
+        return result

bloqade/qasm2/passes/glob.py

@@ -0,0 +1,119 @@
+"""
+Passes that deal with global gates. As of now, only one rewrite pass exists
+which converts global gates to single qubit gates.
+"""
+
+from kirin import ir
+from kirin.rewrite import cse, dce, walk, result
+from kirin.passes.abc import Pass
+from kirin.passes.fold import Fold
+from kirin.rewrite.fixpoint import Fixpoint
+
+from bloqade.analysis import address
+from bloqade.qasm2.rewrite import GlobalToUOpRule, GlobalToParallelRule
+
+
+class GlobalToUOP(Pass):
+    """Pass to convert Global gates into single gates.
+
+    This pass rewrites the global unitary gate from the `qasm2.glob` dialect into multiple
+    single gates in the `qasm2.uop` dialect, bringing the program closer to
+    conforming to standard QASM2 syntax.
+
+
+    ## Usage Examples
+    ```
+    # Define kernel
+    @qasm2.extended
+    def main():
+        q1 = qasm2.qreg(1)
+        q2 = qasm2.qreg(2)
+
+        theta = 1.3
+        phi = 1.1
+        lam = 1.2
+
+        qasm2.glob.u(theta=theta, phi=phi, lam=lam, registers=[q1, q2])
+
+    GlobalToUOP(dialects=main.dialects)(main)
+
+    # Run rewrite
+    GlobalToUOP(main.dialects)(main)
+    ```
+
+    The `qasm2.glob.u` statement has been rewritten to individual gates:
+
+    ```
+    qasm2.uop.u(q1[0], theta, phi, lam)
+    qasm2.uop.u(q2[0], theta, phi, lam)
+    qasm2.uop.u(q2[1], theta, phi, lam)
+    ```
+    """
+
+    def generate_rule(self, mt: ir.Method) -> GlobalToUOpRule:
+        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        return GlobalToUOpRule(frame.entries)
+
+    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+        rewriter = walk.Walk(self.generate_rule(mt))
+        result = rewriter.rewrite(mt.code)
+
+        result = walk.Walk(dce.DeadCodeElimination()).rewrite(mt.code)
+        result = Fixpoint(walk.Walk(rule=cse.CommonSubexpressionElimination())).rewrite(
+            mt.code
+        )
+
+        # do fold again to get proper hint for inserted const
+        result = Fold(mt.dialects)(mt)
+        return result
+
+
+class GlobalToParallel(Pass):
+    """Pass to convert Global gates into parallel gates.
+
+    This pass rewrites the global unitary gate from the `qasm2.glob` dialect into multiple
+    parallel gates in the `qasm2.parallel` dialect.
+
+
+    ## Usage Examples
+    ```
+    # Define kernel
+    @qasm2.extended
+    def main():
+        q1 = qasm2.qreg(1)
+        q2 = qasm2.qreg(2)
+
+        theta = 1.3
+        phi = 1.1
+        lam = 1.2
+
+        qasm2.glob.u(theta=theta, phi=phi, lam=lam, registers=[q1, q2])
+
+    GlobalToParallel(dialects=main.dialects)(main)
+
+    # Run rewrite
+    GlobalToParallel(main.dialects)(main)
+    ```
+
+    The `qasm2.glob.u` statement has been rewritten to individual gates:
+
+    ```
+    qasm2.parallel.u(theta=theta, phi=phi, lam=lam, qargs=[q1[0], q2[0], q2[1]])
+    ```
+    """
+
+    def generate_rule(self, mt: ir.Method) -> GlobalToParallelRule:
+        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        return GlobalToParallelRule(frame.entries)
+
+    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+        rewriter = walk.Walk(self.generate_rule(mt))
+        result = rewriter.rewrite(mt.code)
+
+        result = walk.Walk(dce.DeadCodeElimination()).rewrite(mt.code)
+        result = Fixpoint(walk.Walk(rule=cse.CommonSubexpressionElimination())).rewrite(
+            mt.code
+        )
+        # do fold again to get proper hint
+        result = Fold(mt.dialects)(mt)
+        return result