bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/dialects/core/_emit.py

@@ -0,0 +1,68 @@
+from kirin import interp
+
+from bloqade.qasm2.parse import ast
+from bloqade.qasm2.emit.main import EmitQASM2Main, EmitQASM2Frame
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="emit.qasm2.main")
+class Core(interp.MethodTable):
+
+    @interp.impl(stmts.CRegNew)
+    def emit_creg_new(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: stmts.CRegNew
+    ):
+        n_bits = emit.assert_node(ast.Number, frame.get(stmt.n_bits))
+        # check if its int first, because Int.is_integer() is added for >=3.12
+        assert isinstance(n_bits.value, int), "expected integer"
+        name = emit.ssa_id[stmt.result]
+        frame.body.append(ast.CReg(name=name, size=int(n_bits.value)))
+        return (ast.Name(name),)
+
+    @interp.impl(stmts.QRegNew)
+    def emit_qreg_new(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: stmts.QRegNew
+    ):
+        n_bits = emit.assert_node(ast.Number, frame.get(stmt.n_qubits))
+        assert isinstance(n_bits.value, int), "expected integer"
+        name = emit.ssa_id[stmt.result]
+        frame.body.append(ast.QReg(name=name, size=int(n_bits.value)))
+        return (ast.Name(name),)
+
+    @interp.impl(stmts.Reset)
+    def emit_reset(self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: stmts.Reset):
+        qarg = emit.assert_node((ast.Bit, ast.Name), frame.get(stmt.qarg))
+        frame.body.append(ast.Reset(qarg=qarg))
+        return ()
+
+    @interp.impl(stmts.Measure)
+    def emit_measure(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: stmts.Measure
+    ):
+        qarg = emit.assert_node((ast.Bit, ast.Name), frame.get(stmt.qarg))
+        carg = emit.assert_node((ast.Bit, ast.Name), frame.get(stmt.carg))
+        frame.body.append(ast.Measure(qarg=qarg, carg=carg))
+        return ()
+
+    @interp.impl(stmts.CRegEq)
+    def emit_creg_eq(
+        self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: stmts.CRegEq
+    ):
+        lhs = emit.assert_node(ast.Expr, frame.get(stmt.lhs))
+        rhs = emit.assert_node(ast.Expr, frame.get(stmt.rhs))
+        return (ast.Cmp(lhs=lhs, rhs=rhs),)
+
+    @interp.impl(stmts.CRegGet)
+    @interp.impl(stmts.QRegGet)
+    def emit_qreg_get(
+        self,
+        emit: EmitQASM2Main,
+        frame: EmitQASM2Frame,
+        stmt: stmts.QRegGet | stmts.CRegGet,
+    ):
+        reg = emit.assert_node(ast.Name, frame.get(stmt.reg))
+        idx = emit.assert_node(ast.Number, frame.get(stmt.idx))
+        assert isinstance(idx.value, int)
+        return (ast.Bit(reg, int(idx.value)),)

bloqade/qasm2/dialects/core/_typeinfer.py

@@ -0,0 +1,23 @@
+from kirin import types, interp
+from kirin.analysis import TypeInference
+from kirin.dialects import py
+
+from bloqade.qasm2.types import CRegType, QRegType, QubitType
+
+from ._dialect import dialect
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(py.indexing.GetItem, QRegType, types.Int)
+    def getitem_qreg(
+        self, infer: TypeInference, frame: interp.Frame, node: py.indexing.GetItem
+    ):
+        return (QubitType,)
+
+    @interp.impl(py.indexing.GetItem, CRegType, types.Int)
+    def getitem_creg(
+        self, infer: TypeInference, frame: interp.Frame, node: py.indexing.GetItem
+    ):
+        return (QubitType,)

bloqade/qasm2/dialects/core/address.py

@@ -0,0 +1,38 @@
+from kirin import interp
+
+from bloqade.analysis.address import (
+    Address,
+    NotQubit,
+    AddressReg,
+    AddressQubit,
+    AddressAnalysis,
+)
+
+from .stmts import QRegGet, QRegNew
+from ._dialect import dialect
+
+
+@dialect.register(key="qubit.address")
+class AddressMethodTable(interp.MethodTable):
+
+    @interp.impl(QRegNew)
+    def new(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame[Address],
+        stmt: QRegNew,
+    ):
+        n_qubits = interp.get_const_value(int, stmt.n_qubits)
+        addr = AddressReg(range(interp.next_address, interp.next_address + n_qubits))
+        interp.next_address += n_qubits
+        return (addr,)
+
+    @interp.impl(QRegGet)
+    def get(self, interp: AddressAnalysis, frame: interp.Frame[Address], stmt: QRegGet):
+        addr = frame.get(stmt.reg)
+        pos = interp.get_const_value(int, stmt.idx)
+        if isinstance(addr, AddressReg):
+            global_idx = addr.data[pos]
+            return (AddressQubit(global_idx),)
+        else:  # this is not reachable
+            return (NotQubit(),)

bloqade/qasm2/dialects/core/stmts.py

@@ -0,0 +1,94 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+
+from bloqade.qasm2.types import BitType, CRegType, QRegType, QubitType
+
+from ._dialect import dialect
+
+
+@statement(dialect=dialect)
+class QRegNew(ir.Statement):
+    """Create a new quantum register."""
+
+    name = "qreg.new"
+    traits = frozenset({lowering.FromPythonCall()})
+    n_qubits: ir.SSAValue = info.argument(types.Int)
+    """n_qubits: The number of qubits in the register."""
+    result: ir.ResultValue = info.result(QRegType)
+    """A new quantum register with n_qubits set to |0>."""
+
+
+@statement(dialect=dialect)
+class CRegNew(ir.Statement):
+    """Create a new classical register."""
+
+    name = "creg.new"
+    traits = frozenset({lowering.FromPythonCall()})
+    n_bits: ir.SSAValue = info.argument(types.Int)
+    """n_bits (Int): The number of bits in the register."""
+    result: ir.ResultValue = info.result(CRegType)
+    """result (CReg): The new classical register with all bits set to 0."""
+
+
+@statement(dialect=dialect)
+class Reset(ir.Statement):
+    """Reset a qubit to the |0> state."""
+
+    name = "reset"
+    traits = frozenset({lowering.FromPythonCall()})
+    qarg: ir.SSAValue = info.argument(QubitType)
+    """qarg (Qubit): The qubit to reset."""
+
+
+@statement(dialect=dialect)
+class Measure(ir.Statement):
+    """Measure a qubit and store the result in a bit."""
+
+    name = "measure"
+    traits = frozenset({lowering.FromPythonCall()})
+    qarg: ir.SSAValue = info.argument(QubitType)
+    """qarg (Qubit): The qubit to measure."""
+    carg: ir.SSAValue = info.argument(BitType)
+    """carg (Bit): The bit to store the result in."""
+
+
+@statement(dialect=dialect)
+class CRegEq(ir.Statement):
+    """Check if two classical registers are equal."""
+
+    name = "eq"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+    lhs: ir.SSAValue = info.argument(types.Int | CRegType | BitType)
+    """lhs (CReg): The first register."""
+    rhs: ir.SSAValue = info.argument(types.Int | CRegType | BitType)
+    """rhs (CReg): The second register."""
+    result: ir.ResultValue = info.result(types.Bool)
+    """result (bool): True if the registers are equal, False otherwise."""
+
+
+@statement(dialect=dialect)
+class QRegGet(ir.Statement):
+    """Get a qubit from a quantum register."""
+
+    name = "qreg.get"
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    reg: ir.SSAValue = info.argument(QRegType)
+    """reg (QReg): The quantum register."""
+    idx: ir.SSAValue = info.argument(types.Int)
+    """idx (Int): The index of the qubit in the register."""
+    result: ir.ResultValue = info.result(QubitType)
+    """result (Qubit): The qubit at position `idx`."""
+
+
+@statement(dialect=dialect)
+class CRegGet(ir.Statement):
+    """Get a bit from a classical register."""
+
+    name = "creg.get"
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    reg: ir.SSAValue = info.argument(CRegType)
+    """reg (CReg): The classical register."""
+    idx: ir.SSAValue = info.argument(types.Int)
+    """idx (Int): The index of the bit in the register."""
+    result: ir.ResultValue = info.result(BitType)
+    """result (Bit): The bit at position `idx`."""

bloqade/qasm2/dialects/expr/__init__.py

@@ -0,0 +1,3 @@
+from . import _emit as _emit, _interp as _interp, _from_python as _from_python
+from .stmts import *  # noqa: F403
+from ._dialect import dialect as dialect

bloqade/qasm2/dialects/expr/_dialect.py

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

bloqade/qasm2/dialects/expr/_emit.py

@@ -0,0 +1,103 @@
+from typing import Literal
+
+from kirin import interp
+from kirin.emit.exceptions import EmitError
+
+from bloqade.qasm2.parse import ast
+from bloqade.qasm2.types import QubitType
+from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register(key="emit.qasm2.gate")
+class EmitExpr(interp.MethodTable):
+
+    @interp.impl(stmts.GateFunction)
+    def emit_func(
+        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.GateFunction
+    ):
+
+        cparams, qparams = [], []
+        for arg in stmt.body.blocks[0].args[1:]:
+            name = frame.get(arg)
+            if not isinstance(name, ast.Name):
+                raise EmitError("expected ast.Name")
+            if arg.type.is_subseteq(QubitType):
+                qparams.append(name.id)
+            else:
+                cparams.append(name.id)
+        emit.run_ssacfg_region(frame, stmt.body)
+        emit.output = ast.Gate(
+            name=stmt.sym_name,
+            cparams=cparams,
+            qparams=qparams,
+            body=frame.body,
+        )
+        return ()
+
+    @interp.impl(stmts.ConstInt)
+    @interp.impl(stmts.ConstFloat)
+    def emit_const_int(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: stmts.ConstInt | stmts.ConstFloat,
+    ):
+        return (ast.Number(stmt.value),)
+
+    @interp.impl(stmts.ConstPI)
+    def emit_const_pi(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: stmts.ConstPI,
+    ):
+        return (ast.Pi(),)
+
+    @interp.impl(stmts.Neg)
+    def emit_neg(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.Neg):
+        arg = emit.assert_node(ast.Expr, frame.get(stmt.value))
+        return (ast.UnaryOp("-", arg),)
+
+    @interp.impl(stmts.Sin)
+    @interp.impl(stmts.Cos)
+    @interp.impl(stmts.Tan)
+    @interp.impl(stmts.Exp)
+    @interp.impl(stmts.Log)
+    @interp.impl(stmts.Sqrt)
+    def emit_sin(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
+        arg = emit.assert_node(ast.Expr, frame.get(stmt.value))
+        return (ast.Call(stmt.name, [arg]),)
+
+    def emit_binop(
+        self,
+        sym: Literal["+", "-", "*", "/", "^"],
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt,
+    ):
+        lhs = emit.assert_node(ast.Expr, frame.get(stmt.lhs))
+        rhs = emit.assert_node(ast.Expr, frame.get(stmt.rhs))
+        return (ast.BinOp(sym, lhs, rhs),)
+
+    @interp.impl(stmts.Add)
+    def emit_add(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.Add):
+        return self.emit_binop("+", emit, frame, stmt)
+
+    @interp.impl(stmts.Sub)
+    def emit_sub(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.Add):
+        return self.emit_binop("-", emit, frame, stmt)
+
+    @interp.impl(stmts.Mul)
+    def emit_mul(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.Add):
+        return self.emit_binop("*", emit, frame, stmt)
+
+    @interp.impl(stmts.Div)
+    def emit_div(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.Add):
+        return self.emit_binop("/", emit, frame, stmt)
+
+    @interp.impl(stmts.Pow)
+    def emit_pow(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.Add):
+        return self.emit_binop("^", emit, frame, stmt)

bloqade/qasm2/dialects/expr/_from_python.py

@@ -0,0 +1,86 @@
+import ast
+
+from kirin import ir, types, lowering
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class QASMUopLowering(lowering.FromPythonAST):
+
+    def lower_Name(self, state: lowering.State, node: ast.Name):
+        name = node.id
+        if isinstance(node.ctx, ast.Load):
+            value = state.current_frame.get(name)
+            if value is None:
+                raise lowering.BuildError(f"{name} is not defined")
+            return value
+        elif isinstance(node.ctx, ast.Store):
+            raise lowering.BuildError("unhandled store operation")
+        else:  # Del
+            raise lowering.BuildError("unhandled del operation")
+
+    def lower_Assign(self, state: lowering.State, node: ast.Assign):
+        # NOTE: QASM only expects one value on right hand side
+        rhs = state.lower(node.value).expect_one()
+        current_frame = state.current_frame
+        match node:
+            case ast.Assign(targets=[ast.Name(lhs_name, ast.Store())], value=_):
+                rhs.name = lhs_name
+                current_frame.defs[lhs_name] = rhs
+            case _:
+                target_syntax = ", ".join(
+                    ast.unparse(target) for target in node.targets
+                )
+                raise lowering.BuildError(f"unsupported target syntax {target_syntax}")
+
+    def lower_Expr(self, state: lowering.State, node: ast.Expr):
+        return state.parent.visit(state, node.value)
+
+    def lower_Constant(self, state: lowering.State, node: ast.Constant):
+        if isinstance(node.value, int):
+            stmt = stmts.ConstInt(value=node.value)
+            return state.current_frame.push(stmt)
+        elif isinstance(node.value, float):
+            stmt = stmts.ConstFloat(value=node.value)
+            return state.current_frame.push(stmt)
+        else:
+            raise lowering.BuildError(
+                f"unsupported QASM 2.0 constant type {type(node.value)}"
+            )
+
+    def lower_BinOp(self, state: lowering.State, node: ast.BinOp):
+        lhs = state.lower(node.left).expect_one()
+        rhs = state.lower(node.right).expect_one()
+        if isinstance(node.op, ast.Add):
+            stmt = stmts.Add(lhs, rhs)
+        elif isinstance(node.op, ast.Sub):
+            stmt = stmts.Sub(lhs, rhs)
+        elif isinstance(node.op, ast.Mult):
+            stmt = stmts.Mul(lhs, rhs)
+        elif isinstance(node.op, ast.Div):
+            stmt = stmts.Div(lhs, rhs)
+        elif isinstance(node.op, ast.Pow):
+            stmt = stmts.Pow(lhs, rhs)
+        else:
+            raise lowering.BuildError(f"unsupported QASM 2.0 binop {node.op}")
+        stmt.result.type = self.__promote_binop_type(lhs, rhs)
+        return state.current_frame.push(stmt)
+
+    def __promote_binop_type(
+        self, lhs: ir.SSAValue, rhs: ir.SSAValue
+    ) -> types.TypeAttribute:
+        if lhs.type.is_subseteq(types.Float) or rhs.type.is_subseteq(types.Float):
+            return types.Float
+        return types.Int
+
+    def lower_UnaryOp(self, state: lowering.State, node: ast.UnaryOp):
+        if isinstance(node.op, ast.USub):
+            value = state.lower(node.operand).expect_one()
+            stmt = stmts.Neg(value)
+            return state.current_frame.push(stmt)
+        elif isinstance(node.op, ast.UAdd):
+            return state.lower(node.operand).expect_one()
+        else:
+            raise lowering.BuildError(f"unsupported QASM 2.0 unaryop {node.op}")

bloqade/qasm2/dialects/expr/_interp.py

@@ -0,0 +1,75 @@
+import math
+from typing import Union
+
+from kirin.interp import Frame, Interpreter, MethodTable, impl
+
+from . import stmts
+from ._dialect import dialect
+
+
+@dialect.register
+class Qasm2UopInterpreter(MethodTable):
+    name = "qasm2.uop"
+    dialect = dialect
+
+    @impl(stmts.ConstFloat)
+    @impl(stmts.ConstInt)
+    def new_const(
+        self,
+        interp: Interpreter,
+        frame: Frame,
+        stmt: Union[stmts.ConstFloat, stmts.ConstInt],
+    ):
+        return (stmt.value,)
+
+    @impl(stmts.ConstPI)
+    def new_const_pi(self, interp: Interpreter, frame: Frame, stmt: stmts.ConstPI):
+        return (3.141592653589793,)
+
+    @impl(stmts.Add)
+    def add(self, interp: Interpreter, frame: Frame, stmt: stmts.Add):
+        return (frame.get(stmt.lhs) + frame.get(stmt.rhs),)
+
+    @impl(stmts.Sub)
+    def sub(self, interp: Interpreter, frame: Frame, stmt: stmts.Sub):
+        return (frame.get(stmt.lhs) - frame.get(stmt.rhs),)
+
+    @impl(stmts.Mul)
+    def mul(self, interp: Interpreter, frame: Frame, stmt: stmts.Mul):
+        return (frame.get(stmt.lhs) * frame.get(stmt.rhs),)
+
+    @impl(stmts.Div)
+    def div(self, interp: Interpreter, frame: Frame, stmt: stmts.Div):
+        return (frame.get(stmt.lhs) / frame.get(stmt.rhs),)
+
+    @impl(stmts.Pow)
+    def pow(self, interp: Interpreter, frame: Frame, stmt: stmts.Pow):
+        return (frame.get(stmt.lhs) ** frame.get(stmt.rhs),)
+
+    @impl(stmts.Neg)
+    def neg(self, interp: Interpreter, frame: Frame, stmt: stmts.Neg):
+        return (-frame.get(stmt.value),)
+
+    @impl(stmts.Sqrt)
+    def sqrt(self, interp: Interpreter, frame: Frame, stmt: stmts.Sqrt):
+        return (math.sqrt(frame.get(stmt.value)),)
+
+    @impl(stmts.Sin)
+    def sin(self, interp: Interpreter, frame: Frame, stmt: stmts.Sin):
+        return (math.sin(frame.get(stmt.value)),)
+
+    @impl(stmts.Cos)
+    def cos(self, interp: Interpreter, frame: Frame, stmt: stmts.Cos):
+        return (math.cos(frame.get(stmt.value)),)
+
+    @impl(stmts.Tan)
+    def tan(self, interp: Interpreter, frame: Frame, stmt: stmts.Tan):
+        return (math.tan(frame.get(stmt.value)),)
+
+    @impl(stmts.Log)
+    def log(self, interp: Interpreter, frame: Frame, stmt: stmts.Log):
+        return (math.log(frame.get(stmt.value)),)
+
+    @impl(stmts.Exp)
+    def exp(self, interp: Interpreter, frame: Frame, stmt: stmts.Exp):
+        return (math.exp(frame.get(stmt.value)),)