bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/dialects/expr/stmts.py

@@ -0,0 +1,262 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.print.printer import Printer
+from kirin.dialects.func.attrs import Signature
+
+from ._dialect import dialect
+
+
+class GateFuncOpCallableInterface(ir.CallableStmtInterface["GateFunction"]):
+
+    @classmethod
+    def get_callable_region(cls, stmt: "GateFunction") -> ir.Region:
+        return stmt.body
+
+
+@statement(dialect=dialect)
+class GateFunction(ir.Statement):
+    """Special Function for qasm2 gate subroutine."""
+
+    name = "gate.func"
+    traits = frozenset(
+        {
+            ir.IsolatedFromAbove(),
+            ir.SymbolOpInterface(),
+            ir.HasSignature(),
+            GateFuncOpCallableInterface(),
+        }
+    )
+    sym_name: str = info.attribute()
+    signature: Signature = info.attribute()
+    body: ir.Region = info.region(multi=True)
+
+    def print_impl(self, printer: Printer) -> None:
+        with printer.rich(style="red"):
+            printer.plain_print(self.name + " ")
+
+        with printer.rich(style="cyan"):
+            printer.plain_print(self.sym_name)
+
+        self.signature.print_impl(printer)
+        printer.plain_print(" ")
+        self.body.print_impl(printer)
+
+        with printer.rich(style="black"):
+            printer.plain_print(f" // gate.func {self.sym_name}")
+
+
+@statement(dialect=dialect)
+class ConstInt(ir.Statement):
+    """IR Statement representing a constant integer value."""
+
+    name = "constant.int"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})
+    value: int = info.attribute(types.Int)
+    """value (int): The constant integer value."""
+    result: ir.ResultValue = info.result(types.Int)
+    """result (Int): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class ConstFloat(ir.Statement):
+    """IR Statement representing a constant float value."""
+
+    name = "constant.float"
+    traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})
+    value: float = info.attribute(types.Float)
+    """value (float): The constant float value."""
+    result: ir.ResultValue = info.result(types.Float)
+    """result (Float): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print(repr(self.value))
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+@statement(dialect=dialect)
+class ConstPI(ir.Statement):
+    """The constant value of PI."""
+
+    # this is marked as constant but not pure.
+    name = "constant.pi"
+    traits = frozenset({ir.ConstantLike(), lowering.FromPythonCall()})
+    result: ir.ResultValue = info.result(types.Float)
+    """result (ConstPI): The result value."""
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.plain_print("PI")
+        with printer.rich(style="comment"):
+            printer.plain_print(" : ")
+            printer.print(self.result.type)
+
+
+# QASM 2.0 arithmetic operations
+PyNum = types.Union(types.Int, types.Float)
+
+
+@statement(dialect=dialect)
+class Neg(ir.Statement):
+    """Negate a number."""
+
+    name = "neg"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to negate."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (Union[int, float]): The negated number."""
+
+
+@statement(dialect=dialect)
+class Sin(ir.Statement):
+    """Take the sine of a number."""
+
+    name = "sin"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to take the sine of."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (float): The sine of the number."""
+
+
+@statement(dialect=dialect)
+class Cos(ir.Statement):
+    """Take the cosine of a number."""
+
+    name = "cos"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to take the cosine of."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (float): The cosine of the number."""
+
+
+@statement(dialect=dialect)
+class Tan(ir.Statement):
+    """Take the tangent of a number."""
+
+    name = "tan"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to take the tangent of."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (float): The tangent of the number."""
+
+
+@statement(dialect=dialect)
+class Exp(ir.Statement):
+    """Take the exponential of a number."""
+
+    name = "exp"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to take the exponential of."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (float): The exponential of the number."""
+
+
+@statement(dialect=dialect)
+class Log(ir.Statement):
+    """Take the natural log of a number."""
+
+    name = "ln"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to take the natural log of."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (float): The natural log of the number."""
+
+
+@statement(dialect=dialect)
+class Sqrt(ir.Statement):
+    """Take the square root of a number."""
+
+    name = "sqrt"
+    traits = frozenset({lowering.FromPythonCall()})
+    value: ir.SSAValue = info.argument(PyNum)
+    """value (Union[int, float]): The number to take the square root of."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (float): The square root of the number."""
+
+
+@statement(dialect=dialect)
+class Add(ir.Statement):
+    """Add two numbers."""
+
+    name = "add"
+    traits = frozenset({lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(PyNum)
+    """lhs (Union[int, float]): The left-hand side of the addition."""
+    rhs: ir.SSAValue = info.argument(PyNum)
+    """rhs (Union[int, float]): The right-hand side of the addition."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (Union[int, float]): The result of the addition."""
+
+
+@statement(dialect=dialect)
+class Sub(ir.Statement):
+    """Subtract two numbers."""
+
+    name = "sub"
+    traits = frozenset({lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(PyNum)
+    """lhs (Union[int, float]): The left-hand side of the subtraction."""
+    rhs: ir.SSAValue = info.argument(PyNum)
+    """rhs (Union[int, float]): The right-hand side of the subtraction."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (Union[int, float]): The result of the subtraction."""
+
+
+@statement(dialect=dialect)
+class Mul(ir.Statement):
+    """Multiply two numbers."""
+
+    name = "mul"
+    traits = frozenset({lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(PyNum)
+    """lhs (Union[int, float]): The left-hand side of the multiplication."""
+    rhs: ir.SSAValue = info.argument(PyNum)
+    """rhs (Union[int, float]): The right-hand side of the multiplication."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (Union[int, float]): The result of the multiplication."""
+
+
+@statement(dialect=dialect)
+class Pow(ir.Statement):
+    """Take the power of a number."""
+
+    name = "pow"
+    traits = frozenset({lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(PyNum)
+    """lhs (Union[int, float]): The base."""
+    rhs: ir.SSAValue = info.argument(PyNum)
+    """rhs (Union[int, float]): The exponent."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (Union[int, float]): The result of the power operation."""
+
+
+@statement(dialect=dialect)
+class Div(ir.Statement):
+    """Divide two numbers."""
+
+    name = "div"
+    traits = frozenset({lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(PyNum)
+    """lhs (Union[int, float]): The numerator."""
+    rhs: ir.SSAValue = info.argument(PyNum)
+    """rhs (Union[int, float]): The denominator."""
+    result: ir.ResultValue = info.result(PyNum)
+    """result (Union[int, float]): The result of the division."""

bloqade/qasm2/dialects/glob.py

@@ -0,0 +1,45 @@
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.qasm2.parse import ast
+from bloqade.qasm2.types import QRegType
+from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
+from bloqade.squin.analysis.schedule import DagScheduleAnalysis
+
+dialect = ir.Dialect("qasm2.glob")
+
+
+@statement(dialect=dialect)
+class UGate(ir.Statement):
+    name = "ugate"
+    traits = frozenset({lowering.FromPythonCall()})
+    registers: ir.SSAValue = info.argument(ilist.IListType[QRegType])
+    theta: ir.SSAValue = info.argument(types.Float)
+    phi: ir.SSAValue = info.argument(types.Float)
+    lam: ir.SSAValue = info.argument(types.Float)
+
+
+@dialect.register(key="qasm2.schedule.dag")
+class Glob(interp.MethodTable):
+    @interp.impl(UGate)
+    def ugate(self, interp: DagScheduleAnalysis, frame: interp.Frame, stmt: UGate):
+        interp.update_dag(stmt, [stmt.registers])
+        return ()
+
+
+@dialect.register(key="emit.qasm2.gate")
+class GlobEmit(interp.MethodTable):
+    @interp.impl(UGate)
+    def ugate(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: UGate):
+        registers = [
+            emit.assert_node(ast.Name, reg)
+            for reg in frame.get_casted(stmt.registers, ilist.IList)
+        ]
+        theta = emit.assert_node(ast.Expr, frame.get(stmt.theta))
+        phi = emit.assert_node(ast.Expr, frame.get(stmt.phi))
+        lam = emit.assert_node(ast.Expr, frame.get(stmt.lam))
+        frame.body.append(
+            ast.GlobUGate(theta=theta, phi=phi, lam=lam, registers=registers)
+        )
+        return ()

bloqade/qasm2/dialects/indexing.py

@@ -0,0 +1,64 @@
+"""This dialect provides the indexing syntax in Python lowering
+for QASM2 dialects. The dialect itself does not contain new statements.
+
+Using this dialect will be conflict with Python semantics provided by
+`kirin.dialects.py.binop` and `kirin.dialects.py.indexing` dialects.
+"""
+
+import ast
+
+from kirin import ir, types, lowering
+
+from bloqade.qasm2.types import BitType, CRegType, QRegType, QubitType
+from bloqade.qasm2.dialects import core
+
+dialect = ir.Dialect("qasm2.indexing")
+
+
+@dialect.register
+class QASMCoreLowering(lowering.FromPythonAST):
+    def lower_Compare(self, state: lowering.State, node: ast.Compare):
+        lhs = state.lower(node.left).expect_one()
+        if len(node.ops) != 1:
+            raise lowering.BuildError(
+                "only one comparison operator and == is supported for qasm2 lowering"
+            )
+        rhs = state.lower(node.comparators[0]).expect_one()
+        if isinstance(node.ops[0], ast.Eq):
+            stmt = core.CRegEq(lhs, rhs)
+        else:
+            raise lowering.BuildError(
+                f"unsupported comparison operator {node.ops[0]} only Eq is supported."
+            )
+
+        return state.current_frame.push(stmt)
+
+    def lower_Subscript(self, state: lowering.State, node: ast.Subscript):
+        value = state.lower(node.value).expect_one()
+        index = state.lower(node.slice).expect_one()
+
+        if not index.type.is_subseteq(types.Int):
+            raise lowering.BuildError(
+                f"unsupported subscript index type {index.type},"
+                " only integer indices are supported in QASM 2.0"
+            )
+
+        if not isinstance(node.ctx, ast.Load):
+            raise lowering.BuildError(
+                f"unsupported subscript context {node.ctx},"
+                " cannot write to subscript in QASM 2.0"
+            )
+
+        if value.type.is_subseteq(QRegType):
+            stmt = core.QRegGet(reg=value, idx=index)
+            stmt.result.type = QubitType
+        elif value.type.is_subseteq(CRegType):
+            stmt = core.CRegGet(reg=value, idx=index)
+            stmt.result.type = BitType
+        else:
+            raise lowering.BuildError(
+                f"unsupported subscript value type {value.type},"
+                " only QReg and CReg are supported in QASM 2.0"
+            )
+
+        return state.current_frame.push(stmt)

bloqade/qasm2/dialects/inline.py

@@ -0,0 +1,76 @@
+"""Inline QASM dialect.
+
+This dialect allows users to use QASM string as part of a `@qasm2.main` kernel.
+"""
+
+import ast
+import textwrap
+from dataclasses import dataclass
+
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.print import Printer
+
+dialect = ir.Dialect("qasm2.inline")
+
+
+@dataclass(frozen=True)
+class InlineQASMLowering(lowering.FromPythonCall):
+
+    def lower(
+        self, stmt: type, state: lowering.State, node: ast.Call
+    ) -> lowering.Result:
+        from bloqade.qasm2.parse import loads
+        from bloqade.qasm2.parse.lowering import QASM2
+
+        if len(node.args) != 1 or node.keywords:
+            raise lowering.BuildError("InlineQASM takes 1 positional argument")
+        text = node.args[0]
+        # 1. string literal
+        if isinstance(text, ast.Constant) and isinstance(text.value, str):
+            value = text.value
+        elif isinstance(text, ast.Name) and isinstance(text.ctx, ast.Load):
+            value = state.get_global(text).expect(str)
+        else:
+            raise lowering.BuildError(
+                "InlineQASM takes a string literal or global string"
+            )
+
+        from kirin.dialects import ilist
+
+        from bloqade.qasm2.groups import main
+        from bloqade.qasm2.dialects import glob, noise, parallel
+
+        raw = textwrap.dedent(value)
+        qasm_lowering = QASM2(main.union([ilist, glob, noise, parallel]))
+        region = qasm_lowering.run(loads(raw))
+        for qasm_stmt in region.blocks[0].stmts:
+            qasm_stmt.detach()
+            state.current_frame.push(qasm_stmt)
+
+        for block in region.blocks:
+            for qasm_stmt in block.stmts:
+                qasm_stmt.detach()
+                state.current_frame.push(qasm_stmt)
+            state.current_frame.jump_next_block()
+
+
+# NOTE: this is a dummy statement that won't appear in IR.
+# TODO: maybe we should save the string in IR then rewrite?
+#       what would be the use case?
+@statement(dialect=dialect)
+class InlineQASM(ir.Statement):
+    name = "text"
+    traits = frozenset({InlineQASMLowering()})
+    text: str = info.attribute(types.String)
+
+    def __init__(self, text: str) -> None:
+        super().__init__(attributes={"text": ir.PyAttr(text)})
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print('"""')
+        for line in self.text.splitlines():
+            printer.plain_print(line)
+            printer.print_newline()
+        printer.plain_print('"""')

bloqade/qasm2/dialects/noise.py

@@ -0,0 +1,16 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from bloqade.qasm2.types import QubitType
+
+dialect = ir.Dialect("qasm2.noise")
+
+
+@statement(dialect=dialect)
+class Pauli1(ir.Statement):
+    name = "pauli_1"
+    traits = frozenset({lowering.FromPythonCall()})
+    px: ir.SSAValue = info.argument(types.Float)
+    py: ir.SSAValue = info.argument(types.Float)
+    pz: ir.SSAValue = info.argument(types.Float)
+    qarg: ir.SSAValue = info.argument(QubitType)

bloqade/qasm2/dialects/parallel.py

@@ -0,0 +1,110 @@
+from typing import Any
+
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.analysis import ForwardFrame
+from kirin.dialects import ilist
+
+from bloqade.qasm2.parse import ast
+from bloqade.qasm2.types import QubitType
+from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
+from bloqade.squin.analysis.schedule import DagScheduleAnalysis
+
+dialect = ir.Dialect("qasm2.parallel")
+
+N = types.TypeVar("N")
+
+
+@statement(dialect=dialect)
+class CZ(ir.Statement):
+    name = "cz"
+    traits = frozenset({lowering.FromPythonCall()})
+    ctrls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+    qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+
+
+@statement(dialect=dialect)
+class UGate(ir.Statement):
+    name = "u"
+    traits = frozenset({lowering.FromPythonCall()})
+    qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    theta: ir.SSAValue = info.argument(types.Float)
+    phi: ir.SSAValue = info.argument(types.Float)
+    lam: ir.SSAValue = info.argument(types.Float)
+
+
+@statement(dialect=dialect)
+class RZ(ir.Statement):
+    name = "rz"
+    traits = frozenset({lowering.FromPythonCall()})
+    qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    theta: ir.SSAValue = info.argument(types.Float)
+
+
+@dialect.register(key="emit.qasm2.gate")
+class Parallel(interp.MethodTable):
+
+    def _emit_parallel_qargs(
+        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, qargs: ir.SSAValue
+    ):
+        qargs_: ilist.IList[ast.Node, Any] = frame.get(qargs)  # type: ignore
+        return [(emit.assert_node((ast.Name, ast.Bit), qarg),) for qarg in qargs_]
+
+    @interp.impl(UGate)
+    def ugate(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: UGate):
+        qargs = self._emit_parallel_qargs(emit, frame, stmt.qargs)
+        theta = emit.assert_node(ast.Expr, frame.get(stmt.theta))
+        phi = emit.assert_node(ast.Expr, frame.get(stmt.phi))
+        lam = emit.assert_node(ast.Expr, frame.get(stmt.lam))
+        frame.body.append(
+            ast.ParaU3Gate(
+                theta=theta, phi=phi, lam=lam, qargs=ast.ParallelQArgs(qargs=qargs)
+            )
+        )
+        return ()
+
+    @interp.impl(RZ)
+    def rz(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: RZ):
+        qargs = self._emit_parallel_qargs(emit, frame, stmt.qargs)
+        theta = emit.assert_node(ast.Expr, frame.get(stmt.theta))
+        frame.body.append(
+            ast.ParaRZGate(theta=theta, qargs=ast.ParallelQArgs(qargs=qargs))
+        )
+        return ()
+
+    @interp.impl(CZ)
+    def cz(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: CZ):
+        ctrls = self._emit_parallel_qargs(emit, frame, stmt.ctrls)
+        qargs = self._emit_parallel_qargs(emit, frame, stmt.qargs)
+        frame.body.append(
+            ast.ParaCZGate(
+                qargs=ast.ParallelQArgs(
+                    qargs=[ctrl + qarg for ctrl, qarg in zip(ctrls, qargs)]
+                )
+            )
+        )
+        return ()
+
+
+@dialect.register(key="qasm2.schedule.dag")
+class ParallelDag(interp.MethodTable):
+
+    @interp.impl(CZ)
+    def parallel_cz(self, interp: DagScheduleAnalysis, frame: ForwardFrame, stmt: CZ):
+        interp.update_dag(stmt, [stmt.qargs, stmt.ctrls])
+        return ()
+
+    @interp.impl(UGate)
+    def parallel_ugate(
+        self,
+        interp: DagScheduleAnalysis,
+        frame: ForwardFrame,
+        stmt: UGate,
+    ):
+        interp.update_dag(stmt, [stmt.qargs])
+        return ()
+
+    @interp.impl(RZ)
+    def parallel_rz(self, interp: DagScheduleAnalysis, frame: ForwardFrame, stmt: RZ):
+        interp.update_dag(stmt, [stmt.qargs])
+        return ()

bloqade/qasm2/dialects/uop/__init__.py

@@ -0,0 +1,4 @@
+from . import _emit as _emit, stmts as stmts
+from .stmts import *  # noqa: F403
+from ._dialect import dialect as dialect
+from .schedule import *  # noqa: F403

bloqade/qasm2/dialects/uop/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("qasm2.uop")