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

bloqade/qasm2/rewrite/parallel_to_uop.py

@@ -21,16 +21,10 @@ class ParallelToUOpRule(abc.RewriteRule):
 
     def get_qubit_ssa(self, ilist_ref: ir.SSAValue) -> Optional[List[ir.SSAValue]]:
         addr = self.address_analysis.get(ilist_ref)
-        if not isinstance(addr, address.AddressTuple):
+        if not isinstance(addr, address.AddressReg):
             return None
 
-        ids = []
-        for ele in addr.data:
-            if not isinstance(ele, address.AddressQubit):
-                return None
-
-            ids.append(ele.data)
-
+        ids = addr.data
         return [self.id_map[ele] for ele in ids]
 
     def rewrite_cz(self, node: ir.Statement):

bloqade/qasm2/rewrite/register.py

@@ -2,7 +2,7 @@ from kirin import ir
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.qasm2.dialects import core
+from bloqade.qasm2.dialects import core, expr
 
 
 class RaiseRegisterRule(RewriteRule):
@@ -26,7 +26,7 @@ class RaiseRegisterRule(RewriteRule):
         n_qubits_ref = node.n_qubits
 
         n_qubits = n_qubits_ref.owner
-        if isinstance(n_qubits, py.Constant):
+        if isinstance(n_qubits, py.Constant | expr.ConstInt):
             # case where the n_qubits comes from a constant
             new_n_qubits = n_qubits.from_stmt(n_qubits)
             new_n_qubits.insert_before(first_stmt)

bloqade/qasm2/rewrite/uop_to_parallel.py

@@ -8,7 +8,7 @@ from kirin.rewrite.abc import RewriteRule, RewriteResult
 from kirin.analysis.const import lattice
 
 from bloqade.analysis import address
-from bloqade.qasm2.dialects import uop, core, parallel
+from bloqade.qasm2.dialects import uop, core, expr, parallel
 from bloqade.squin.analysis.schedule import StmtDag
 
 
@@ -66,7 +66,7 @@ class SimpleMergePolicy(MergePolicyABC):
             assert isinstance(hint1, lattice.Result) and isinstance(
                 hint2, lattice.Result
             )
-            return hint1.is_equal(hint2)
+            return hint1.is_structurally_equal(hint2)
         else:
             return False
 
@@ -194,6 +194,8 @@ class SimpleMergePolicy(MergePolicyABC):
                     new_qubits.append(new_qubit.result)
                 case core.QRegGet(
                     reg=reg, idx=ir.ResultValue(stmt=py.Constant() as idx)
+                ) | core.QRegGet(
+                    reg=reg, idx=ir.ResultValue(stmt=expr.ConstInt() as idx)
                 ):
                     (new_idx := idx.from_stmt(idx)).insert_before(node)
                     (

bloqade/qbraid/lowering.py

@@ -320,5 +320,6 @@ class Lowering:
         self.block_list.append(const_pi)
         turns = self.lower_number(2 * value)
         mul = qasm2.expr.Mul(const_pi.result, turns)
+        mul.result.type = types.Float
         self.block_list.append(mul)
         return mul.result

bloqade/qbraid/schema.py

@@ -238,13 +238,13 @@ class NoiseModel(BaseModel, Generic[ErrorModelType], extra="forbid"):
             str: The decompiled circuit from hardware execution.
 
         """
-        from bloqade.noise import native
         from bloqade.qasm2.emit import QASM2
         from bloqade.qasm2.passes import glob, parallel
+        from bloqade.qasm2.rewrite.noise import remove_noise
 
         mt = self.lower_noise_model("method")
 
-        native.RemoveNoisePass(mt.dialects)(mt)
+        remove_noise.RemoveNoisePass(mt.dialects)(mt)
         parallel.ParallelToUOp(mt.dialects)(mt)
         glob.GlobalToUOP(mt.dialects)(mt)
         return QASM2(qelib1=True).emit_str(mt)

bloqade/qubit/__init__.py

@@ -0,0 +1,12 @@
+from bloqade.types import Qubit as Qubit, QubitType as QubitType
+
+from . import stmts as stmts, analysis as analysis
+from .stdlib import new as new, qalloc as qalloc, broadcast as broadcast
+from ._dialect import dialect as dialect
+from ._prelude import kernel as kernel
+from .stdlib.simple import (
+    reset as reset,
+    measure as measure,
+    get_qubit_id as get_qubit_id,
+    get_measurement_id as get_measurement_id,
+)

bloqade/qubit/_dialect.py

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

bloqade/qubit/_interface.py

@@ -0,0 +1,49 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit, MeasurementResult
+
+from .stmts import New, Reset, Measure, QubitId, MeasurementId
+
+
+@wraps(New)
+def new() -> Qubit:
+    """Create a new qubit.
+
+    Returns:
+        Qubit: A new qubit.
+    """
+    ...
+
+
+N = TypeVar("N", bound=int)
+
+
+@wraps(Measure)
+def measure(qubits: ilist.IList[Qubit, N]) -> ilist.IList[MeasurementResult, N]:
+    """Measure a list of qubits.
+
+    Args:
+        qubits (IList[Qubit, N]): The list of qubits to measure.
+
+    Returns:
+        IList[MeasurementResult, N]: The list containing the results of the measurements.
+            A MeasurementResult can represent both 0 and 1, but also atoms that are lost.
+    """
+    ...
+
+
+@wraps(QubitId)
+def get_qubit_id(qubits: ilist.IList[Qubit, N]) -> ilist.IList[int, N]: ...
+
+
+@wraps(MeasurementId)
+def get_measurement_id(
+    measurements: ilist.IList[MeasurementResult, N],
+) -> ilist.IList[int, N]: ...
+
+
+@wraps(Reset)
+def reset(qubits: ilist.IList[Qubit, Any]) -> None: ...

bloqade/qubit/_prelude.py

@@ -0,0 +1,45 @@
+from typing import Annotated
+
+from kirin import ir
+from kirin.passes import Default
+from kirin.prelude import structural_no_opt
+from typing_extensions import Doc
+
+from . import _dialect as qubit
+
+
+@ir.dialect_group(structural_no_opt.union([qubit]))
+def kernel(self):
+    """Compile to a qubit kernel"""
+
+    def run_pass(
+        mt,
+        *,
+        verify: Annotated[
+            bool, Doc("run `verify` before running passes, default is `True`")
+        ] = True,
+        typeinfer: Annotated[
+            bool,
+            Doc(
+                "run type inference and apply the inferred type to IR, default `False`"
+            ),
+        ] = False,
+        fold: Annotated[bool, Doc("run folding passes, default is `True`")] = True,
+        aggressive: Annotated[
+            bool, Doc("run aggressive folding passes if `fold=True`")
+        ] = False,
+        no_raise: Annotated[
+            bool, Doc("do not raise exception during analysis, default is `True`")
+        ] = True,
+    ) -> None:
+        default_pass = Default(
+            self,
+            verify=verify,
+            fold=fold,
+            aggressive=aggressive,
+            typeinfer=typeinfer,
+            no_raise=no_raise,
+        )
+        default_pass.fixpoint(mt)
+
+    return run_pass

bloqade/qubit/analysis/__init__.py

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

bloqade/qubit/analysis/address_impl.py

@@ -0,0 +1,40 @@
+from kirin import interp
+from kirin.analysis import ForwardFrame
+
+from bloqade.analysis.address.lattice import (
+    Address,
+    AddressQubit,
+)
+from bloqade.analysis.address.analysis import AddressAnalysis
+
+from .. import stmts
+from .._dialect import dialect
+
+# Address lattice elements we can work with:
+## NotQubit (bottom), AnyAddress (top)
+
+## AddressTuple -> data: tuple[Address, ...]
+### Recursive type, could contain itself or other variants
+### This pops up in cases where you can have an IList/Tuple
+### That contains elements that could be other Address types
+
+## AddressReg -> data: Sequence[int]
+### specific to creation of a register of qubits
+
+## AddressQubit -> data: int
+### Base qubit address type
+
+
+@dialect.register(key="qubit.address")
+class SquinQubitMethodTable(interp.MethodTable):
+
+    @interp.impl(stmts.New)
+    def new_qubit(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: stmts.New,
+    ):
+        addr = AddressQubit(interp_.next_address)
+        interp_.next_address += 1
+        return (addr,)

bloqade/qubit/stdlib/__init__.py

@@ -0,0 +1,2 @@
+from . import simple as simple, broadcast as broadcast
+from ._new import new as new, qalloc as qalloc

bloqade/qubit/stdlib/_new.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from kirin.dialects import ilist
+
+from .. import _interface as qubit
+from .._prelude import kernel
+
+
+@kernel(typeinfer=True)
+def new() -> qubit.Qubit:
+    """Allocate a single new qubit
+
+    Returns:
+        (Qubit): The newly allocated qubit.
+    """
+    return qubit.new()
+
+
+# NOTE: this is a special case, that doesn't use the usual simple / broadcast semantics.
+@kernel(typeinfer=True)
+def qalloc(n_qubits: int) -> ilist.IList[qubit.Qubit, Any]:
+    """Allocate a new list of qubits.
+
+    Args:
+        n_qubits(int): The number of qubits to create.
+
+    Returns:
+        (ilist.IList[Qubit, n_qubits]) A list of qubits.
+    """
+
+    def _new(qid: int) -> qubit.Qubit:
+        return qubit.new()
+
+    return ilist.map(_new, ilist.range(n_qubits))

bloqade/qubit/stdlib/broadcast.py

@@ -0,0 +1,62 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit, MeasurementResult
+
+from .. import _interface as _qubit
+from .._prelude import kernel
+
+N = TypeVar("N", bound=int)
+
+
+@kernel
+def reset(qubits: ilist.IList[Qubit, Any]) -> None:
+    """
+    Reset a list of qubits to the zero state.
+
+    Args:
+        qubits (IList[Qubit, Any]): The list of qubits to reset.
+    """
+    _qubit.reset(qubits)
+
+
+@kernel
+def measure(qubits: ilist.IList[Qubit, N]) -> ilist.IList[MeasurementResult, N]:
+    """Measure a list of qubits.
+
+    Args:
+        qubits (IList[Qubit, N]): The list of qubits to measure.
+
+    Returns:
+        IList[MeasurementResult, N]: The list containing the results of the measurements.
+            A MeasurementResult can represent both 0 and 1 as well as atom loss.
+    """
+    return _qubit.measure(qubits)
+
+
+@kernel
+def get_qubit_id(qubits: ilist.IList[Qubit, N]) -> ilist.IList[int, N]:
+    """Get the global, unique ID of each qubit in the list.
+
+    Args:
+        qubits (IList[Qubit, N]): The list of qubits of which you want the ID.
+
+    Returns:
+        qubit_ids (IList[int, N]): The list of global, unique IDs of the qubits.
+    """
+    return _qubit.get_qubit_id(qubits)
+
+
+@kernel
+def get_measurement_id(
+    measurements: ilist.IList[MeasurementResult, N],
+) -> ilist.IList[int, N]:
+    """Get the global, unique ID of each of the measurement results in the list.
+
+    Args:
+        measurements (IList[MeasurementResult, N]): The previously taken measurement of which you want to know the ID.
+    Returns:
+        measurement_ids (IList[int, N]): The list of global, unique IDs of the measurements.
+    """
+    return _qubit.get_measurement_id(measurements)

bloqade/qubit/stdlib/simple.py

@@ -0,0 +1,59 @@
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit, MeasurementResult
+
+from . import broadcast
+from .._prelude import kernel
+
+
+@kernel
+def reset(qubit: Qubit) -> None:
+    """
+    Reset a qubit to the zero state.
+
+    Args:
+        qubit (Qubit): The list qubit to reset.
+    """
+    return broadcast.reset(ilist.IList([qubit]))
+
+
+@kernel
+def measure(qubit: Qubit) -> MeasurementResult:
+    """Measure a qubit.
+
+    Args:
+        qubit (Qubit): The qubit to measure.
+
+    Returns:
+        MeasurementResult: The result of the measurement.
+            A MeasurementResult can represent both 0 and 1, but also atoms that are lost.
+    """
+    measurement_results = broadcast.measure(ilist.IList([qubit]))
+    return measurement_results[0]
+
+
+@kernel
+def get_qubit_id(qubit: Qubit) -> int:
+    """Get the global, unique ID of the qubit.
+
+    Args:
+        qubit (Qubit): The qubit of which you want the ID.
+
+    Returns:
+        qubit_id (int): The global, unique ID of the qubit.
+    """
+    ids = broadcast.get_qubit_id(ilist.IList([qubit]))
+    return ids[0]
+
+
+@kernel
+def get_measurement_id(measurement: MeasurementResult) -> int:
+    """Get the global, unique ID of the measurement result.
+
+    Args:
+        measurement (MeasurementResult): The previously taken measurement of which you want to know the ID.
+    Returns:
+        measurement_id (int): The global, unique ID of the measurement.
+    """
+    ids = broadcast.get_measurement_id(ilist.IList([measurement]))
+    return ids[0]

bloqade/qubit/stmts.py

@@ -0,0 +1,60 @@
+from kirin import ir, types, interp, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.types import QubitType, MeasurementResultType
+
+from ._dialect import dialect
+
+
+@statement(dialect=dialect)
+class New(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    result: ir.ResultValue = info.result(QubitType)
+
+
+Len = types.TypeVar("Len", bound=types.Int)
+
+
+@statement(dialect=dialect)
+class Measure(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, Len])
+    result: ir.ResultValue = info.result(ilist.IListType[MeasurementResultType, Len])
+
+
+@statement(dialect=dialect)
+class QubitId(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, Len])
+    result: ir.ResultValue = info.result(ilist.IListType[types.Int, Len])
+
+
+@statement(dialect=dialect)
+class MeasurementId(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    measurements: ir.SSAValue = info.argument(
+        ilist.IListType[MeasurementResultType, Len]
+    )
+    result: ir.ResultValue = info.result(ilist.IListType[types.Int, Len])
+
+
+@statement(dialect=dialect)
+class Reset(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+# TODO: investigate why this is needed to get type inference to be correct.
+@dialect.register(key="typeinfer")
+class __TypeInfer(interp.MethodTable):
+    @interp.impl(Measure)
+    def measure_list(self, _interp, frame: interp.AbstractFrame, stmt: Measure):
+        qubit_type = frame.get(stmt.qubits)
+
+        if isinstance(qubit_type, types.Generic):
+            len_type = qubit_type.vars[1]
+        else:
+            len_type = types.Any
+
+        return (ilist.IListType[MeasurementResultType, len_type],)

bloqade/rewrite/passes/__init__.py

@@ -1 +1,7 @@
+from .callgraph import (
+    CallGraphPass as CallGraphPass,
+    ReplaceMethods as ReplaceMethods,
+    UpdateDialectsOnCallGraph as UpdateDialectsOnCallGraph,
+)
+from .aggressive_unroll import AggressiveUnroll as AggressiveUnroll
 from .canonicalize_ilist import CanonicalizeIList as CanonicalizeIList

bloqade/rewrite/passes/aggressive_unroll.py

@@ -0,0 +1,103 @@
+from typing import Callable
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass, HintConst, TypeInfer
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Inline,
+    Fixpoint,
+    Call2Invoke,
+    ConstantFold,
+    CFGCompactify,
+    InlineGetItem,
+    InlineGetField,
+    DeadCodeElimination,
+    CommonSubexpressionElimination,
+)
+from kirin.dialects import scf, ilist
+from kirin.ir.method import Method
+from kirin.rewrite.abc import RewriteResult
+from kirin.passes.aggressive import UnrollScf
+
+from .canonicalize_ilist import CanonicalizeIList
+
+
+@dataclass
+class Fold(Pass):
+    hint_const: HintConst = field(init=False)
+
+    def __post_init__(self):
+        self.hint_const = HintConst(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        result = self.hint_const.unsafe_run(mt).join(result)
+        rule = Chain(
+            ConstantFold(),
+            Call2Invoke(),
+            InlineGetField(),
+            InlineGetItem(),
+            ilist.rewrite.InlineGetItem(),
+            ilist.rewrite.FlattenAdd(),
+            ilist.rewrite.HintLen(),
+            DeadCodeElimination(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+        return result
+
+
+@dataclass
+class AggressiveUnroll(Pass):
+    """A pass to unroll structured control flow"""
+
+    additional_inline_heuristic: Callable[[ir.Statement], bool] = lambda node: True
+
+    fold: Fold = field(init=False)
+    typeinfer: TypeInfer = field(init=False)
+    scf_unroll: UnrollScf = field(init=False)
+    canonicalize_ilist: CanonicalizeIList = field(init=False)
+
+    def __post_init__(self):
+        self.fold = Fold(self.dialects, no_raise=self.no_raise)
+        self.typeinfer = TypeInfer(self.dialects, no_raise=self.no_raise)
+        self.scf_unroll = UnrollScf(self.dialects, no_raise=self.no_raise)
+        self.canonicalize_ilist = CanonicalizeIList(
+            self.dialects, no_raise=self.no_raise
+        )
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        result = self.fold.unsafe_run(mt).join(result)
+        result = self.scf_unroll.unsafe_run(mt).join(result)
+        self.typeinfer.unsafe_run(
+            mt
+        )  # Do not join the result of typeinfer or fixpoint will waste time
+        result = (
+            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        result = Walk(Inline(self.inline_heuristic)).rewrite(mt.code).join(result)
+        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
+        result = self.canonicalize_ilist.fixpoint(mt).join(result)
+        rule = Chain(
+            CommonSubexpressionElimination(),
+            DeadCodeElimination(),
+        )
+        result = Walk(rule).rewrite(mt.code).join(result)
+
+        return result
+
+    def inline_heuristic(self, node: ir.Statement) -> bool:
+        """The heuristic to decide whether to inline a function call or not.
+        inside loops and if-else, only inline simple functions, i.e.
+        functions with a single block
+        """
+        return not isinstance(
+            node.parent_stmt, (scf.For, scf.IfElse)
+        ) and self.additional_inline_heuristic(
+            node
+        )  # always inline calls outside of loops and if-else