bloqade-circuit 0.6.4__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/rewrite/passes/callgraph.py

@@ -0,0 +1,116 @@
+from dataclasses import field, dataclass
+
+from kirin import ir, passes, rewrite
+from kirin.analysis import CallGraph
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.func.stmts import Invoke
+
+
+@dataclass
+class ReplaceMethods(RewriteRule):
+    new_symbols: dict[ir.Method, ir.Method]
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if (
+            not isinstance(node, Invoke)
+            or (new_callee := self.new_symbols.get(node.callee)) is None
+        ):
+            return RewriteResult()
+
+        node.replace_by(
+            Invoke(
+                inputs=node.inputs,
+                callee=new_callee,
+                purity=node.purity,
+            )
+        )
+
+        return RewriteResult(has_done_something=True)
+
+
+@dataclass
+class UpdateDialectsOnCallGraph(passes.Pass):
+    """Update All dialects on the call graph to a new set of dialects given to this pass.
+
+    Usage:
+        pass_ = UpdateDialectsOnCallGraph(rule=rule, dialects=new_dialects)
+        pass_(some_method)
+
+    Note: This pass does not update the dialects of the input method, but copies
+    all other methods invoked within it before updating their dialects.
+
+    """
+
+    fold_pass: passes.Fold = field(init=False)
+
+    def __post_init__(self):
+        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        mt_map = {}
+
+        cg = CallGraph(mt)
+
+        all_methods = set(sum(map(tuple, cg.defs.values()), ()))
+        for original_mt in all_methods:
+            if original_mt is mt:
+                new_mt = original_mt
+            else:
+                new_mt = original_mt.similar(self.dialects)
+            mt_map[original_mt] = new_mt
+
+        result = RewriteResult()
+
+        for _, new_mt in mt_map.items():
+            result = (
+                rewrite.Walk(ReplaceMethods(mt_map)).rewrite(new_mt.code).join(result)
+            )
+            self.fold_pass(new_mt)
+
+        return result
+
+
+@dataclass
+class CallGraphPass(passes.Pass):
+    """Copy all functions in the call graph and apply a rule to each of them.
+
+
+    Usage:
+        rule = Walk(SomeRewriteRule())
+        pass_ = CallGraphPass(rule=rule, dialects=...)
+        pass_(some_method)
+
+    Note: This pass modifies the input method in place, but copies
+    all methods invoked within it before applying the rule to them.
+
+    """
+
+    rule: RewriteRule
+    """The rule to apply to each function in the call graph."""
+
+    fold_pass: passes.Fold = field(init=False)
+
+    def __post_init__(self):
+        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        result = RewriteResult()
+        mt_map = {}
+
+        cg = CallGraph(mt)
+
+        all_methods = set(cg.edges.keys())
+        for original_mt in all_methods:
+            if original_mt is mt:
+                new_mt = original_mt
+            else:
+                new_mt = original_mt.similar()
+            result = self.rule.rewrite(new_mt.code).join(result)
+            mt_map[original_mt] = new_mt
+
+        if result.has_done_something:
+            for _, new_mt in mt_map.items():
+                rewrite.Walk(ReplaceMethods(mt_map)).rewrite(new_mt.code)
+                self.fold_pass(new_mt)
+
+        return result

bloqade/rewrite/passes/canonicalize_ilist.py

@@ -1,28 +1,34 @@
-from dataclasses import dataclass
+from dataclasses import field, dataclass
 
-from kirin import ir
-from kirin.passes import Pass
+from kirin import ir, passes
 from kirin.rewrite import (
     Walk,
     Chain,
     Fixpoint,
+    DeadCodeElimination,
 )
-from kirin.analysis import const
-
-from ..rules.flatten_ilist import FlattenAddOpIList
-from ..rules.inline_getitem_ilist import InlineGetItemFromIList
+from kirin.dialects.ilist import rewrite
 
 
 @dataclass
-class CanonicalizeIList(Pass):
+class CanonicalizeIList(passes.Pass):
 
-    def unsafe_run(self, mt: ir.Method):
+    fold_pass: passes.Fold = field(init=False)
 
-        cp_result_frame, _ = const.Propagate(dialects=mt.dialects).run_analysis(mt)
+    def __post_init__(self):
+        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
 
-        return Fixpoint(
-            Chain(
-                Walk(InlineGetItemFromIList(constprop_result=cp_result_frame.entries)),
-                Walk(FlattenAddOpIList()),
+    def unsafe_run(self, mt: ir.Method):
+        result = Fixpoint(
+            Walk(
+                Chain(
+                    rewrite.InlineGetItem(),
+                    rewrite.FlattenAdd(),
+                    rewrite.HintLen(),
+                    DeadCodeElimination(),
+                )
             )
         ).rewrite(mt.code)
+
+        result = self.fold_pass(mt).join(result)
+        return result

bloqade/rewrite/rules/split_ifs.py

@@ -46,9 +46,13 @@ class SplitIfStmts(RewriteRule):
         if not isinstance(node, scf.IfElse):
             return RewriteResult()
 
+        # NOTE: only empty else bodies are allowed in valid QASM2
+        if not self._has_empty_else(node):
+            return RewriteResult()
+
         *stmts, yield_or_return = node.then_body.stmts()
 
-        if len(stmts) == 1:
+        if len(stmts) <= 1:
             return RewriteResult()
 
         is_yield = isinstance(yield_or_return, scf.Yield)
@@ -71,3 +75,16 @@ class SplitIfStmts(RewriteRule):
         node.delete()
 
         return RewriteResult(has_done_something=True)
+
+    def _has_empty_else(self, node: scf.IfElse) -> bool:
+        else_stmts = list(node.else_body.stmts())
+        if len(else_stmts) > 1:
+            return False
+
+        if len(else_stmts) == 0:
+            return True
+
+        if not isinstance(else_stmts[0], scf.Yield):
+            return False
+
+        return len(else_stmts[0].values) == 0

bloqade/squin/__init__.py

@@ -1,19 +1,52 @@
 from . import (
-    op as op,
-    wire as wire,
+    gate as gate,
     noise as noise,
-    qubit as qubit,
     analysis as analysis,
-    lowering as lowering,
-    _typeinfer as _typeinfer,
 )
-from .groups import wired as wired, kernel as kernel
+from .. import qubit as qubit, annotate as annotate
+from ..qubit import (
+    reset as reset,
+    qalloc as qalloc,
+    measure as measure,
+    get_qubit_id as get_qubit_id,
+    get_measurement_id as get_measurement_id,
+)
+from .groups import kernel as kernel
+from ..annotate import set_detector as set_detector, set_observable as set_observable
+from .stdlib.simple import (
+    h as h,
+    s as s,
+    t as t,
+    x as x,
+    y as y,
+    z as z,
+    cx as cx,
+    cy as cy,
+    cz as cz,
+    rx as rx,
+    ry as ry,
+    rz as rz,
+    u3 as u3,
+    s_adj as s_adj,
+    shift as shift,
+    t_adj as t_adj,
+    sqrt_x as sqrt_x,
+    sqrt_y as sqrt_y,
+    sqrt_z as sqrt_z,
+    bit_flip as bit_flip,
+    depolarize as depolarize,
+    qubit_loss as qubit_loss,
+    sqrt_x_adj as sqrt_x_adj,
+    sqrt_y_adj as sqrt_y_adj,
+    sqrt_z_adj as sqrt_z_adj,
+    depolarize2 as depolarize2,
+    correlated_qubit_loss as correlated_qubit_loss,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
+)
 
-try:
-    # NOTE: make sure optional cirq dependency is installed
-    import cirq as cirq_package  # noqa: F401
-except ImportError:
-    pass
-else:
-    from . import cirq as cirq
-    from .cirq import load_circuit as load_circuit
+# NOTE: it's important to keep these imports here since they import squin.kernel
+# we skip isort here
+from .stdlib import (  # isort: skip
+    broadcast as broadcast,
+)

bloqade/squin/analysis/__init__.py

@@ -1 +0,0 @@
-from . import address_impl as address_impl

bloqade/squin/analysis/schedule.py

@@ -185,18 +185,17 @@ class DagScheduleAnalysis(Forward[GateSchedule]):
         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 method_self(self, method: ir.Method) -> GateSchedule:
+        return self.lattice.bottom()
 
-    def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
-        if stmt.has_trait(ir.IsTerminator):
+    def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
+        if node.has_trait(ir.IsTerminator):
             assert (
-                stmt.parent_block is not None
+                node.parent_block is not None
             ), "Terminator statement has no parent block"
-            self.push_current_dag(stmt.parent_block)
+            self.push_current_dag(node.parent_block)
 
-        return tuple(self.lattice.top() for _ in stmt.results)
+        return tuple(self.lattice.top() for _ in node.results)
 
     def _update_dag(self, stmt: ir.Statement, addr: address.Address):
         if isinstance(addr, address.AddressQubit):
@@ -210,8 +209,8 @@ class DagScheduleAnalysis(Forward[GateSchedule]):
                 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:
+        elif isinstance(addr, address.AddressReg):
+            for sub_addr in addr.qubits:
                 self._update_dag(stmt, sub_addr)
 
     def update_dag(self, stmt: ir.Statement, args: Sequence[ir.SSAValue]):
@@ -226,7 +225,7 @@ class DagScheduleAnalysis(Forward[GateSchedule]):
         if args is None:
             args = tuple(self.lattice.top() for _ in mt.args)
 
-        self.run(mt, args, kwargs)
+        self.run(mt)
         return self.stmt_dags
 
 

bloqade/squin/gate/__init__.py

@@ -0,0 +1,2 @@
+from . import stmts as stmts
+from ._dialect import dialect as dialect

bloqade/squin/gate/_dialect.py

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

bloqade/squin/gate/_interface.py

@@ -0,0 +1,98 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit
+
+from .stmts import (
+    CX,
+    CY,
+    CZ,
+    U3,
+    H,
+    S,
+    T,
+    X,
+    Y,
+    Z,
+    Rx,
+    Ry,
+    Rz,
+    SqrtX,
+    SqrtY,
+)
+
+
+@wraps(X)
+def x(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Y)
+def y(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Z)
+def z(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(H)
+def h(qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(T)
+def t(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(S)
+def s(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(SqrtX)
+def sqrt_x(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(SqrtY)
+def sqrt_y(qubits: ilist.IList[Qubit, Any], *, adjoint: bool) -> None: ...
+
+
+@wraps(Rx)
+def rx(angle: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Ry)
+def ry(angle: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(Rz)
+def rz(angle: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+Len = TypeVar("Len", bound=int)
+
+
+@wraps(CX)
+def cx(
+    controls: ilist.IList[Qubit, Len],
+    targets: ilist.IList[Qubit, Len],
+) -> None: ...
+
+
+@wraps(CY)
+def cy(
+    controls: ilist.IList[Qubit, Len],
+    targets: ilist.IList[Qubit, Len],
+) -> None: ...
+
+
+@wraps(CZ)
+def cz(
+    controls: ilist.IList[Qubit, Len],
+    targets: ilist.IList[Qubit, Len],
+) -> None: ...
+
+
+@wraps(U3)
+def u3(
+    theta: float, phi: float, lam: float, qubits: ilist.IList[Qubit, Any]
+) -> None: ...

bloqade/squin/gate/stmts.py

@@ -0,0 +1,125 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.types import QubitType
+
+from ._dialect import dialect
+
+
+@statement
+class Gate(ir.Statement):
+    # NOTE: just for easier isinstance checks elsewhere, all gates inherit from this class
+    pass
+
+
+@statement
+class SingleQubitGate(Gate):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+@statement(dialect=dialect)
+class X(SingleQubitGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Y(SingleQubitGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Z(SingleQubitGate):
+    pass
+
+
+@statement(dialect=dialect)
+class H(SingleQubitGate):
+    pass
+
+
+@statement
+class SingleQubitNonHermitianGate(SingleQubitGate):
+    adjoint: bool = info.attribute(default=False)
+
+
+@statement(dialect=dialect)
+class T(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement(dialect=dialect)
+class S(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement(dialect=dialect)
+class SqrtX(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement(dialect=dialect)
+class SqrtY(SingleQubitNonHermitianGate):
+    pass
+
+
+@statement
+class RotationGate(Gate):
+    # NOTE: don't inherit from SingleQubitGate here so the wrapper doesn't have qubits as first arg
+    traits = frozenset({lowering.FromPythonCall()})
+    angle: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+@statement(dialect=dialect)
+class Rx(RotationGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Ry(RotationGate):
+    pass
+
+
+@statement(dialect=dialect)
+class Rz(RotationGate):
+    pass
+
+
+N = types.TypeVar("N", bound=types.Int)
+
+
+@statement
+class ControlledGate(Gate):
+    traits = frozenset({lowering.FromPythonCall()})
+    controls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+    targets: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+
+
+@statement(dialect=dialect)
+class CX(ControlledGate):
+    name = "cx"
+    pass
+
+
+@statement(dialect=dialect)
+class CY(ControlledGate):
+    name = "cy"
+    pass
+
+
+@statement(dialect=dialect)
+class CZ(ControlledGate):
+    name = "cz"
+    pass
+
+
+@statement(dialect=dialect)
+class U3(Gate):
+    # NOTE: don't inherit from SingleQubitGate here so the wrapper doesn't have qubits as first arg
+    traits = frozenset({lowering.FromPythonCall()})
+    theta: ir.SSAValue = info.argument(types.Float)
+    phi: ir.SSAValue = info.argument(types.Float)
+    lam: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])

bloqade/squin/groups.py

@@ -1,31 +1,24 @@
 from kirin import ir, passes
 from kirin.prelude import structural_no_opt
-from kirin.rewrite import Walk, Chain
-from kirin.dialects import ilist
+from kirin.dialects import debug, ilist
 
-from . import op, wire, noise, qubit
-from .op.rewrite import PyMultToSquinMult
-from .rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
+from . import gate, noise
+from .. import qubit, annotate
 
 
-@ir.dialect_group(structural_no_opt.union([op, qubit, noise]))
+@ir.dialect_group(structural_no_opt.union([qubit, noise, gate, debug, annotate]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
     ilist_desugar_pass = ilist.IListDesugar(self)
-    desugar_pass = Walk(Chain(MeasureDesugarRule(), ApplyDesugarRule()))
-    py_mult_to_mult_pass = PyMultToSquinMult(self)
 
     def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
         method.verify()
         if fold:
             fold_pass.fixpoint(method)
 
-        py_mult_to_mult_pass(method)
-
         if typeinfer:
-            typeinfer_pass(method)
-            desugar_pass.rewrite(method.code)
+            typeinfer_pass(method)  # infer types before desugaring
 
         ilist_desugar_pass(method)
 
@@ -34,13 +27,3 @@ def kernel(self):
             method.verify_type()
 
     return run_pass
-
-
-@ir.dialect_group(structural_no_opt.union([op, wire, noise]))
-def wired(self):
-    py_mult_to_mult_pass = PyMultToSquinMult(self)
-
-    def run_pass(method):
-        py_mult_to_mult_pass(method)
-
-    return run_pass

bloqade/squin/noise/__init__.py

@@ -1,11 +1,2 @@
-from . import stmts as stmts
+from . import stmts as stmts, _interface as _interface
 from ._dialect import dialect as dialect
-from ._wrapper import (
-    pp_error as pp_error,
-    depolarize as depolarize,
-    qubit_loss as qubit_loss,
-    depolarize2 as depolarize2,
-    pauli_error as pauli_error,
-    two_qubit_pauli_channel as two_qubit_pauli_channel,
-    single_qubit_pauli_channel as single_qubit_pauli_channel,
-)

bloqade/squin/noise/_dialect.py

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

bloqade/squin/noise/_interface.py

@@ -0,0 +1,45 @@
+from typing import Any, Literal, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit
+
+from . import stmts
+
+
+@wraps(stmts.Depolarize)
+def depolarize(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+N = TypeVar("N", bound=int)
+
+
+@wraps(stmts.Depolarize2)
+def depolarize2(
+    p: float, controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]
+) -> None: ...
+
+
+@wraps(stmts.SingleQubitPauliChannel)
+def single_qubit_pauli_channel(
+    px: float, py: float, pz: float, qubits: ilist.IList[Qubit, Any]
+) -> None: ...
+
+
+@wraps(stmts.TwoQubitPauliChannel)
+def two_qubit_pauli_channel(
+    probabilities: ilist.IList[float, Literal[15]],
+    controls: ilist.IList[Qubit, N],
+    targets: ilist.IList[Qubit, N],
+) -> None: ...
+
+
+@wraps(stmts.QubitLoss)
+def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(stmts.CorrelatedQubitLoss)
+def correlated_qubit_loss(
+    p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any]
+) -> None: ...