bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qbraid/target.py

@@ -0,0 +1,86 @@
+from typing import TYPE_CHECKING, Union, Optional
+
+from kirin import ir
+
+if TYPE_CHECKING:
+    from qbraid import QbraidProvider
+    from qbraid.runtime import QbraidJob
+
+from bloqade.qasm2.emit import QASM2
+
+
+class qBraid:
+    """qBraid target for Bloqade kernels.
+
+    qBraid target that accepts a Bloqade kernel and submits the kernel to the QuEra simulator hosted on qBraid. A `QbraidJob` is obtainable
+    that then lets you query the status of the submitted program on the simulator as well
+    as obtain results.
+
+    """
+
+    def __init__(
+        self,
+        *,
+        allow_parallel: bool = False,
+        allow_global: bool = False,
+        provider: "QbraidProvider",  # inject externally for easier mocking
+        qelib1: bool = True,
+    ) -> None:
+        """Initialize the qBraid 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.
+
+            provider (QbraidProvider):
+                Qbraid-provided object to allow submission of the kernel to the QuEra simulator.
+            qelib1 (bool):
+                Include the `include "qelib1.inc"` line in the resulting QASM2 AST that's
+                submitted to qBraid. Defaults to `True`.
+        """
+
+        self.qelib1 = qelib1
+        self.provider = provider
+        self.allow_parallel = allow_parallel
+        self.allow_global = allow_global
+
+    def emit(
+        self,
+        method: ir.Method,
+        shots: Optional[int] = None,
+        tags: Optional[dict[str, str]] = None,
+    ) -> Union["QbraidJob", list["QbraidJob"]]:
+        """Submit the Bloqade kernel to the QuEra simulator on qBraid.
+
+        Args:
+            method (ir.Method):
+                The kernel to submit to qBraid.
+            shots: (Optional[int]):
+                Number of times to run the kernel. Defaults to None.
+            tags: (Optional[dict[str,str]]):
+                A dictionary of tags to associate with the Job.
+
+        Returns:
+            Union[QbraidJob, list[QbraidJob]]:
+                An object you can query for the status of your submission as well as
+                obtain simulator results from.
+        """
+
+        # Convert method to QASM2 string
+        qasm2_emitter = QASM2(
+            allow_parallel=self.allow_parallel,
+            allow_global=self.allow_global,
+            qelib1=self.qelib1,
+        )
+        qasm2_prog = qasm2_emitter.emit_str(method)
+
+        # Submit the QASM2 string to the qBraid simulator
+        quera_qasm_simulator = self.provider.get_device("quera_qasm_simulator")
+
+        return quera_qasm_simulator.run(qasm2_prog, shots=shots, tags=tags)

bloqade/squin/__init__.py

@@ -0,0 +1,2 @@
+from . import op as op, wire as wire, qubit as qubit
+from .groups import wired as wired, kernel as kernel

bloqade/squin/analysis/nsites/__init__.py

@@ -0,0 +1,8 @@
+# Need this for impl registration to work properly!
+from . import impls as impls
+from .lattice import (
+    NoSites as NoSites,
+    AnySites as AnySites,
+    NumberSites as NumberSites,
+)
+from .analysis import NSitesAnalysis as NSitesAnalysis

bloqade/squin/analysis/nsites/analysis.py

@@ -0,0 +1,52 @@
+# from typing import cast
+
+from kirin import ir
+from kirin.analysis import Forward
+from kirin.analysis.forward import ForwardFrame
+
+from bloqade.squin.op.types import OpType
+from bloqade.squin.op.traits import HasSites, FixedSites
+
+from .lattice import Sites, NoSites, NumberSites
+
+
+class NSitesAnalysis(Forward[Sites]):
+
+    keys = ["op.nsites"]
+    lattice = Sites
+
+    # Take a page from const prop in Kirin,
+    # I can get the data I want from the SizedTrait
+    # and go from there
+
+    ## This gets called before the registry look up
+    def eval_stmt(self, frame: ForwardFrame, stmt: ir.Statement):
+        method = self.lookup_registry(frame, stmt)
+        if method is not None:
+            return method(self, frame, stmt)
+        elif stmt.has_trait(HasSites):
+            has_sites_trait = stmt.get_trait(HasSites)
+            sites = has_sites_trait.get_sites(stmt)
+            return (NumberSites(sites=sites),)
+        elif stmt.has_trait(FixedSites):
+            sites_trait = stmt.get_trait(FixedSites)
+            return (NumberSites(sites=sites_trait.data),)
+        else:
+            return (NoSites(),)
+
+    # For when no implementation is found for the statement
+    def eval_stmt_fallback(
+        self, frame: ForwardFrame[Sites], stmt: ir.Statement
+    ) -> tuple[Sites, ...]:  # some form of Shape will go back into the frame
+        return tuple(
+            (
+                self.lattice.top()
+                if result.type.is_subseteq(OpType)
+                else self.lattice.bottom()
+            )
+            for result in stmt.results
+        )
+
+    def run_method(self, method: ir.Method, args: tuple[Sites, ...]):
+        # NOTE: we do not support dynamic calls here, thus no need to propagate method object
+        return self.run_callable(method.code, (self.lattice.bottom(),) + args)

bloqade/squin/analysis/nsites/impls.py

@@ -0,0 +1,69 @@
+from typing import cast
+
+from kirin import ir, interp
+
+from bloqade.squin import op
+
+from .lattice import (
+    NoSites,
+    NumberSites,
+)
+from .analysis import NSitesAnalysis
+
+
+@op.dialect.register(key="op.nsites")
+class SquinOp(interp.MethodTable):
+
+    @interp.impl(op.stmts.Kron)
+    def kron(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Kron):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        if isinstance(lhs, NumberSites) and isinstance(rhs, NumberSites):
+            new_n_sites = lhs.sites + rhs.sites
+            return (NumberSites(sites=new_n_sites),)
+        else:
+            return (NoSites(),)
+
+    @interp.impl(op.stmts.Mult)
+    def mult(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Mult):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+
+        if isinstance(lhs, NumberSites) and isinstance(rhs, NumberSites):
+            lhs_sites = lhs.sites
+            rhs_sites = rhs.sites
+            # I originally considered throwing an exception here
+            # but Xiu-zhe (Roger) Luo has pointed out it would be
+            # a much better UX to add a type element that
+            # could explicitly indicate the error. The downside
+            # is you'll have some added complexity in the type lattice.
+            if lhs_sites != rhs_sites:
+                return (NoSites(),)
+            else:
+                return (NumberSites(sites=lhs_sites + rhs_sites),)
+        else:
+            return (NoSites(),)
+
+    @interp.impl(op.stmts.Control)
+    def control(
+        self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Control
+    ):
+        op_sites = frame.get(stmt.op)
+
+        if isinstance(op_sites, NumberSites):
+            n_sites = op_sites.sites
+            n_controls_attr = stmt.get_attr_or_prop("n_controls")
+            n_controls = cast(ir.PyAttr[int], n_controls_attr).data
+            return (NumberSites(sites=n_sites + n_controls),)
+        else:
+            return (NoSites(),)
+
+    @interp.impl(op.stmts.Rot)
+    def rot(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Rot):
+        op_sites = frame.get(stmt.axis)
+        return (op_sites,)
+
+    @interp.impl(op.stmts.Scale)
+    def scale(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Scale):
+        op_sites = frame.get(stmt.op)
+        return (op_sites,)

bloqade/squin/analysis/nsites/lattice.py

@@ -0,0 +1,49 @@
+from typing import final
+from dataclasses import dataclass
+
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+
+
+@dataclass
+class Sites(
+    SimpleJoinMixin["Sites"], SimpleMeetMixin["Sites"], BoundedLattice["Sites"]
+):
+    @classmethod
+    def bottom(cls) -> "Sites":
+        return NoSites()
+
+    @classmethod
+    def top(cls) -> "Sites":
+        return AnySites()
+
+
+@final
+@dataclass
+class NoSites(Sites, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: Sites) -> bool:
+        return True
+
+
+@final
+@dataclass
+class AnySites(Sites, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: Sites) -> bool:
+        return isinstance(other, Sites)
+
+
+@final
+@dataclass
+class NumberSites(Sites):
+    sites: int
+
+    def is_subseteq(self, other: Sites) -> bool:
+        if isinstance(other, NumberSites):
+            return self.sites == other.sites
+        return False

bloqade/squin/analysis/schedule.py

@@ -0,0 +1,244 @@
+from typing import Any, Set, Dict, Iterable, Optional, final
+from itertools import chain
+from collections import OrderedDict
+from dataclasses import field, dataclass
+from collections.abc import Sequence
+
+from kirin import ir, graph, interp, idtable
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+from kirin.analysis import Forward, ForwardFrame
+from kirin.dialects import func
+
+from bloqade.analysis import address
+from bloqade.qasm2.parse.print import Printer
+
+
+@dataclass
+class GateSchedule(
+    SimpleJoinMixin["GateSchedule"],
+    SimpleMeetMixin["GateSchedule"],
+    BoundedLattice["GateSchedule"],
+):
+
+    @classmethod
+    def bottom(cls) -> "GateSchedule":
+        return NotQubit()
+
+    @classmethod
+    def top(cls) -> "GateSchedule":
+        return Qubit()
+
+
+@final
+@dataclass
+class NotQubit(GateSchedule, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: GateSchedule) -> bool:
+        return True
+
+
+@final
+@dataclass
+class Qubit(GateSchedule, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: GateSchedule) -> bool:
+        return isinstance(other, Qubit)
+
+
+# Treat global gates as terminators for this analysis, e.g. split block in half.
+
+
+@dataclass(slots=True)
+class StmtDag(graph.Graph[ir.Statement]):
+    id_table: idtable.IdTable[ir.Statement] = field(
+        default_factory=lambda: idtable.IdTable()
+    )
+    stmts: Dict[str, ir.Statement] = field(default_factory=OrderedDict)
+    out_edges: Dict[str, Set[str]] = field(default_factory=OrderedDict)
+    inc_edges: Dict[str, Set[str]] = field(default_factory=OrderedDict)
+    stmt_index: Dict[ir.Statement, int] = field(default_factory=OrderedDict)
+
+    def update_index(self, node: ir.Statement):
+        if node not in self.stmt_index:
+            self.stmt_index[node] = len(self.stmt_index)
+
+    def add_node(self, node: ir.Statement):
+        node_id = self.id_table[node]
+        self.stmts[node_id] = node
+        self.update_index(node)
+        self.out_edges.setdefault(node_id, set())
+        self.inc_edges.setdefault(node_id, set())
+        return node_id
+
+    def add_edge(self, src: ir.Statement, dst: ir.Statement):
+        src_id = self.add_node(src)
+        dst_id = self.add_node(dst)
+
+        self.out_edges[src_id].add(dst_id)
+        self.inc_edges[dst_id].add(src_id)
+
+    def get_parents(self, node: ir.Statement) -> Iterable[ir.Statement]:
+        return (
+            self.stmts[node_id]
+            for node_id in self.inc_edges.get(self.id_table[node], set())
+        )
+
+    def get_children(self, node: ir.Statement) -> Iterable[ir.Statement]:
+        return (
+            self.stmts[node_id]
+            for node_id in self.out_edges.get(self.id_table[node], set())
+        )
+
+    def get_neighbors(self, node: ir.Statement) -> Iterable[ir.Statement]:
+        return chain(self.get_parents(node), self.get_children(node))
+
+    def get_nodes(self) -> Iterable[ir.Statement]:
+        return self.stmts.values()
+
+    def get_edges(self) -> Iterable[tuple[ir.Statement, ir.Statement]]:
+        return (
+            (self.stmts[src], self.stmts[dst])
+            for src, dsts in self.out_edges.items()
+            for dst in dsts
+        )
+
+    def print(
+        self,
+        printer: Optional["Printer"] = None,
+        analysis: dict["ir.SSAValue", Any] | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    def topological_groups(self):
+        """Split the dag into topological groups where each group
+        contains nodes that have no dependencies on each other, but
+        have dependencies on nodes in one or more previous groups.
+
+        Yields:
+            List[str]: A list of node ids in a topological group
+
+
+        Raises:
+            ValueError: If a cyclic dependency is detected
+
+
+        The idea is to yield all nodes with no dependencies, then remove
+        those nodes from the graph repeating until no nodes are left
+        or we reach some upper limit. Worse case is a linear dag,
+        so we can use len(dag.stmts) as the upper limit
+
+        If we reach the limit and there are still nodes left, then we
+        have a cyclic dependency.
+        """
+
+        inc_edges = {k: set(v) for k, v in self.inc_edges.items()}
+
+        check_next = inc_edges.keys()
+
+        for _ in range(len(self.stmts)):
+            if len(inc_edges) == 0:
+                break
+
+            group = [node_id for node_id in check_next if len(inc_edges[node_id]) == 0]
+            yield group
+
+            check_next = set()
+            for n in group:
+                inc_edges.pop(n)
+                for m in self.out_edges[n]:
+                    check_next.add(m)
+                    inc_edges[m].remove(n)
+
+        if inc_edges:
+            raise ValueError("Cyclic dependency detected")
+
+
+@dataclass
+class DagScheduleAnalysis(Forward[GateSchedule]):
+    keys = ["qasm2.schedule.dag"]
+    lattice = GateSchedule
+
+    address_analysis: Dict[ir.SSAValue, address.Address]
+    use_def: Dict[int, ir.Statement] = field(init=False)
+    stmt_dag: StmtDag = field(init=False)
+    stmt_dags: Dict[ir.Block, StmtDag] = field(init=False)
+
+    def initialize(self):
+        self.use_def = {}
+        self.stmt_dag = StmtDag()
+        self.stmt_dags = {}
+        return super().initialize()
+
+    def push_current_dag(self, block: ir.Block):
+        # run when hitting terminator statements
+        assert block not in self.stmt_dags, "Block already in stmt_dags"
+
+        for node in self.use_def.values():
+            self.stmt_dag.add_node(node)
+
+        self.stmt_dags[block] = self.stmt_dag
+        self.stmt_dag = StmtDag()
+        self.use_def = {}
+
+    def run_method(self, method: ir.Method, args: tuple[GateSchedule, ...]):
+        # NOTE: we do not support dynamic calls here, thus no need to propagate method object
+        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+
+    def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
+        if stmt.has_trait(ir.IsTerminator):
+            assert (
+                stmt.parent_block is not None
+            ), "Terminator statement has no parent block"
+            self.push_current_dag(stmt.parent_block)
+
+        return tuple(self.lattice.top() for _ in stmt.results)
+
+    def _update_dag(self, stmt: ir.Statement, addr: address.Address):
+        if isinstance(addr, address.AddressQubit):
+            old_stmt = self.use_def.get(addr.data, None)
+            if old_stmt is not None:
+                self.stmt_dag.add_edge(old_stmt, stmt)
+            self.use_def[addr.data] = stmt
+        elif isinstance(addr, address.AddressReg):
+            for idx in addr.data:
+                old_stmt = self.use_def.get(idx, None)
+                if old_stmt is not None:
+                    self.stmt_dag.add_edge(old_stmt, stmt)
+                self.use_def[idx] = stmt
+        elif isinstance(addr, address.AddressTuple):
+            for sub_addr in addr.data:
+                self._update_dag(stmt, sub_addr)
+
+    def update_dag(self, stmt: ir.Statement, args: Sequence[ir.SSAValue]):
+        self.stmt_dag.add_node(stmt)
+
+        for arg in args:
+            self._update_dag(
+                stmt, self.address_analysis.get(arg, address.Address.bottom())
+            )
+
+    def get_dags(self, mt: ir.Method, args=None, kwargs=None):
+        if args is None:
+            args = tuple(self.lattice.top() for _ in mt.args)
+
+        self.run(mt, args, kwargs).expect()
+        return self.stmt_dags
+
+
+@func.dialect.register(key="qasm2.schedule.dag")
+class FuncImpl(interp.MethodTable):
+    @interp.impl(func.Invoke)
+    @interp.impl(func.Call)
+    def invoke(
+        self,
+        interp: DagScheduleAnalysis,
+        frame: ForwardFrame,
+        stmt: func.Invoke | func.Call,
+    ):
+        interp.update_dag(stmt, stmt.inputs)
+        return tuple(interp.lattice.top() for _ in stmt.results)

bloqade/squin/groups.py

@@ -0,0 +1,38 @@
+from kirin import ir, passes
+from kirin.prelude import structural_no_opt
+from kirin.dialects import ilist
+
+from bloqade.qasm2.rewrite.desugar import IndexingDesugarPass
+
+from . import op, wire, qubit
+
+
+@ir.dialect_group(structural_no_opt.union([op, qubit]))
+def kernel(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(method, *, fold=True, typeinfer=True):
+        method.verify()
+        if fold:
+            fold_pass.fixpoint(method)
+
+        if typeinfer:
+            typeinfer_pass(method)
+        ilist_desugar_pass(method)
+        indexing_desugar_pass(method)
+        if typeinfer:
+            typeinfer_pass(method)  # fix types after desugaring
+            method.code.typecheck()
+
+    return run_pass
+
+
+@ir.dialect_group(structural_no_opt.union([op, wire]))
+def wired(self):
+    def run_pass(method):
+        pass
+
+    return run_pass

bloqade/squin/op/__init__.py

@@ -0,0 +1,132 @@
+from kirin import ir as _ir
+from kirin.prelude import structural_no_opt as _structural_no_opt
+from kirin.lowering import wraps as _wraps
+
+from . import stmts as stmts, types as types
+from .traits import Unitary as Unitary, MaybeUnitary as MaybeUnitary
+from ._dialect import dialect as dialect
+
+
+@_wraps(stmts.Kron)
+def kron(lhs: types.Op, rhs: types.Op, *, is_unitary: bool = False) -> types.Op: ...
+
+
+@_wraps(stmts.Adjoint)
+def adjoint(op: types.Op, *, is_unitary: bool = False) -> types.Op: ...
+
+
+@_wraps(stmts.Control)
+def control(op: types.Op, *, n_controls: int, is_unitary: bool = False) -> types.Op: ...
+
+
+@_wraps(stmts.Identity)
+def identity(*, size: int) -> types.Op: ...
+
+
+@_wraps(stmts.Rot)
+def rot(axis: types.Op, angle: float) -> types.Op: ...
+
+
+@_wraps(stmts.ShiftOp)
+def shift(theta: float) -> types.Op: ...
+
+
+@_wraps(stmts.PhaseOp)
+def phase(theta: float) -> types.Op: ...
+
+
+@_wraps(stmts.X)
+def x() -> types.Op: ...
+
+
+@_wraps(stmts.Y)
+def y() -> types.Op: ...
+
+
+@_wraps(stmts.Z)
+def z() -> types.Op: ...
+
+
+@_wraps(stmts.H)
+def h() -> types.Op: ...
+
+
+@_wraps(stmts.S)
+def s() -> types.Op: ...
+
+
+@_wraps(stmts.T)
+def t() -> types.Op: ...
+
+
+@_wraps(stmts.P0)
+def p0() -> types.Op: ...
+
+
+@_wraps(stmts.P1)
+def p1() -> types.Op: ...
+
+
+@_wraps(stmts.Sn)
+def spin_n() -> types.Op: ...
+
+
+@_wraps(stmts.Sp)
+def spin_p() -> types.Op: ...
+
+
+# stdlibs
+@_ir.dialect_group(_structural_no_opt.add(dialect))
+def op(self):
+    def run_pass(method):
+        pass
+
+    return run_pass
+
+
+@op
+def rx(theta: float) -> types.Op:
+    """Rotation X gate."""
+    return rot(x(), theta)
+
+
+@op
+def ry(theta: float) -> types.Op:
+    """Rotation Y gate."""
+    return rot(y(), theta)
+
+
+@op
+def rz(theta: float) -> types.Op:
+    """Rotation Z gate."""
+    return rot(z(), theta)
+
+
+@op
+def cx() -> types.Op:
+    """Controlled X gate."""
+    return control(x(), n_controls=1)
+
+
+@op
+def cy() -> types.Op:
+    """Controlled Y gate."""
+    return control(y(), n_controls=1)
+
+
+@op
+def cz() -> types.Op:
+    """Control Z gate."""
+    return control(z(), n_controls=1)
+
+
+@op
+def ch() -> types.Op:
+    """Control H gate."""
+    return control(h(), n_controls=1)
+
+
+@op
+def cphase(theta: float) -> types.Op:
+    """Control Phase gate."""
+    return control(phase(theta), n_controls=1)

bloqade/squin/op/_dialect.py

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

bloqade/squin/op/complex.py

@@ -0,0 +1,6 @@
+# Stopgap Measure, squin dialect needs Complex type but
+# this is only available in Kirin 0.15.x
+
+from kirin.ir.attrs.types import PyClass
+
+Complex = PyClass(complex)