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

bloqade/stim/rewrite/util.py

@@ -1,35 +1,8 @@
-from typing import TypeVar
-
-from kirin import ir, interp
-from kirin.analysis import const
+from kirin import ir
 from kirin.dialects import py
-from kirin.rewrite.abc import RewriteResult
 
-from bloqade.squin import op, wire, noise as squin_noise, qubit
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.stim.dialects import gate, noise as stim_noise, collapse
-from bloqade.analysis.address import AddressReg, AddressWire, AddressQubit, AddressTuple
-
-SQUIN_STIM_OP_MAPPING = {
-    op.stmts.X: gate.X,
-    op.stmts.Y: gate.Y,
-    op.stmts.Z: gate.Z,
-    op.stmts.H: gate.H,
-    op.stmts.S: gate.S,
-    op.stmts.SqrtX: gate.SqrtX,
-    op.stmts.SqrtY: gate.SqrtY,
-    op.stmts.Identity: gate.Identity,
-    op.stmts.Reset: collapse.RZ,
-    squin_noise.stmts.QubitLoss: stim_noise.QubitLoss,
-}
-
-# Squin allows creation of control gates where the gate can be any operator,
-# but Stim only supports CX, CY, and CZ as control gates.
-SQUIN_STIM_CONTROL_GATE_MAPPING = {
-    op.stmts.X: gate.CX,
-    op.stmts.Y: gate.CY,
-    op.stmts.Z: gate.CZ,
-}
+from bloqade.analysis.address import AddressReg, AddressQubit
 
 
 def create_and_insert_qubit_idx_stmt(
@@ -46,177 +19,17 @@ def insert_qubit_idx_from_address(
     """
     Extract qubit indices from an AddressAttribute and insert them into the SSA form.
     """
-    address_data = address.address
     qubit_idx_ssas = []
-
-    if isinstance(address_data, AddressTuple):
-        for address_qubit in address_data.data:
-            if not isinstance(address_qubit, AddressQubit):
-                return
+    if isinstance(address_data := address.address, AddressReg):
+        for qubit_idx in address_data.qubits:
             create_and_insert_qubit_idx_stmt(
-                address_qubit.data, stmt_to_insert_before, qubit_idx_ssas
-            )
-    elif isinstance(address_data, AddressReg):
-        for qubit_idx in address_data.data:
-            create_and_insert_qubit_idx_stmt(
-                qubit_idx, stmt_to_insert_before, qubit_idx_ssas
+                qubit_idx.data, stmt_to_insert_before, qubit_idx_ssas
             )
     elif isinstance(address_data, AddressQubit):
         create_and_insert_qubit_idx_stmt(
             address_data.data, stmt_to_insert_before, qubit_idx_ssas
         )
-    elif isinstance(address_data, AddressWire):
-        address_qubit = address_data.origin_qubit
-        create_and_insert_qubit_idx_stmt(
-            address_qubit.data, stmt_to_insert_before, qubit_idx_ssas
-        )
     else:
         return
 
     return tuple(qubit_idx_ssas)
-
-
-def insert_qubit_idx_from_wire_ssa(
-    wire_ssas: tuple[ir.SSAValue, ...], stmt_to_insert_before: ir.Statement
-) -> tuple[ir.SSAValue, ...] | None:
-    """
-    Extract qubit indices from wire SSA values and insert them into the SSA form.
-    """
-    qubit_idx_ssas = []
-    for wire_ssa in wire_ssas:
-        address_attribute = wire_ssa.hints.get("address")
-        if address_attribute is None:
-            return
-        assert isinstance(address_attribute, AddressAttribute)
-        wire_address = address_attribute.address
-        assert isinstance(wire_address, AddressWire)
-        qubit_idx = wire_address.origin_qubit.data
-        qubit_idx_stmt = py.Constant(qubit_idx)
-        qubit_idx_ssas.append(qubit_idx_stmt.result)
-        qubit_idx_stmt.insert_before(stmt_to_insert_before)
-
-    return tuple(qubit_idx_ssas)
-
-
-def insert_qubit_idx_after_apply(
-    stmt: wire.Apply | qubit.Apply | wire.Broadcast | qubit.Broadcast,
-) -> tuple[ir.SSAValue, ...] | None:
-    """
-    Extract qubit indices from Apply or Broadcast statements.
-    """
-    if isinstance(stmt, (qubit.Apply, qubit.Broadcast)):
-        qubits = stmt.qubits
-        address_attribute = qubits.hints.get("address")
-        if address_attribute is None:
-            return
-        assert isinstance(address_attribute, AddressAttribute)
-        return insert_qubit_idx_from_address(
-            address=address_attribute, stmt_to_insert_before=stmt
-        )
-    elif isinstance(stmt, (wire.Apply, wire.Broadcast)):
-        wire_ssas = stmt.inputs
-        return insert_qubit_idx_from_wire_ssa(
-            wire_ssas=wire_ssas, stmt_to_insert_before=stmt
-        )
-
-
-def rewrite_Control(
-    stmt_with_ctrl: qubit.Apply | wire.Apply | qubit.Broadcast | wire.Broadcast,
-) -> RewriteResult:
-    """
-    Handle control gates for Apply and Broadcast statements.
-    """
-    ctrl_op = stmt_with_ctrl.operator.owner
-    assert isinstance(ctrl_op, op.stmts.Control)
-
-    ctrl_op_target_gate = ctrl_op.op.owner
-    assert isinstance(ctrl_op_target_gate, op.stmts.Operator)
-
-    qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt_with_ctrl)
-    if qubit_idx_ssas is None:
-        return RewriteResult()
-
-    # Separate control and target qubits
-    target_qubits = []
-    ctrl_qubits = []
-    for i in range(len(qubit_idx_ssas)):
-        if (i % 2) == 0:
-            ctrl_qubits.append(qubit_idx_ssas[i])
-        else:
-            target_qubits.append(qubit_idx_ssas[i])
-
-    target_qubits = tuple(target_qubits)
-    ctrl_qubits = tuple(ctrl_qubits)
-
-    stim_gate = SQUIN_STIM_CONTROL_GATE_MAPPING.get(type(ctrl_op_target_gate))
-    if stim_gate is None:
-        return RewriteResult()
-
-    stim_stmt = stim_gate(controls=ctrl_qubits, targets=target_qubits)
-
-    if isinstance(stmt_with_ctrl, (wire.Apply, wire.Broadcast)):
-        create_wire_passthrough(stmt_with_ctrl)
-
-    stmt_with_ctrl.replace_by(stim_stmt)
-
-    return RewriteResult(has_done_something=True)
-
-
-def rewrite_QubitLoss(
-    stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
-) -> RewriteResult:
-    """
-    Rewrite QubitLoss statements to Stim's TrivialError.
-    """
-
-    squin_loss_op = stmt.operator.owner
-    assert isinstance(squin_loss_op, squin_noise.stmts.QubitLoss)
-
-    qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt)
-    if qubit_idx_ssas is None:
-        return RewriteResult()
-
-    stim_loss_stmt = stim_noise.QubitLoss(
-        targets=qubit_idx_ssas,
-        probs=(squin_loss_op.p,),
-    )
-
-    if isinstance(stmt, (wire.Apply, wire.Broadcast)):
-        create_wire_passthrough(stmt)
-
-    stmt.replace_by(stim_loss_stmt)
-
-    return RewriteResult(has_done_something=True)
-
-
-def create_wire_passthrough(stmt: wire.Apply | wire.Broadcast) -> None:
-
-    for input_wire, output_wire in zip(stmt.inputs, stmt.results):
-        # have to "reroute" the input of these statements to directly plug in
-        # to subsequent statements, remove dependency on the current statement
-        output_wire.replace_by(input_wire)
-
-
-def is_measure_result_used(
-    stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure,
-) -> bool:
-    """
-    Check if the result of a measure statement is used in the program.
-    """
-    return bool(stmt.result.uses)
-
-
-T = TypeVar("T")
-
-
-def get_const_value(typ: type[T], value: ir.SSAValue) -> T:
-    if isinstance(hint := value.hints.get("const"), const.Value):
-        data = hint.data
-        if isinstance(data, typ):
-            return hint.data
-        raise interp.InterpreterError(
-            f"Expected constant value <type = {typ}>, got {data}"
-        )
-    raise interp.InterpreterError(
-        f"Expected constant value <type = {typ}>, got {value}"
-    )

bloqade/test_utils.py

@@ -25,7 +25,7 @@ def print_diff(node: pprint.Printable, expected_node: pprint.Printable):
 
 def assert_nodes(node: ir.IRNode, expected_node: ir.IRNode):
     try:
-        assert node.is_equal(expected_node)
+        assert node.is_structurally_equal(expected_node)
     except AssertionError as e:
         print_diff(node, expected_node)
         raise e

bloqade/types.py

@@ -22,3 +22,13 @@ class Qubit(ABC):
 
 QubitType = types.PyClass(Qubit)
 """Kirin type for a qubit."""
+
+
+class MeasurementResult:
+    """Runtime representation of the result of a measurement on a qubit."""
+
+    pass
+
+
+MeasurementResultType = types.PyClass(MeasurementResult)
+"""Kirin type for a measurement result."""

bloqade/validation/__init__.py

@@ -0,0 +1,2 @@
+from . import analysis as analysis
+from .kernel_validation import KernelValidation as KernelValidation

bloqade/validation/analysis/__init__.py

@@ -0,0 +1,5 @@
+from . import lattice as lattice
+from .analysis import (
+    ValidationFrame as ValidationFrame,
+    ValidationAnalysis as ValidationAnalysis,
+)

bloqade/validation/analysis/analysis.py

@@ -0,0 +1,41 @@
+from abc import ABC
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.analysis import ForwardExtra, ForwardFrame
+
+from .lattice import ErrorType
+
+
+@dataclass
+class ValidationFrame(ForwardFrame[ErrorType]):
+    # NOTE: cannot be set[Error] since that's not hashable
+    errors: list[ir.ValidationError] = field(default_factory=list)
+    """List of all ecnountered errors.
+
+    Append a `kirin.ir.ValidationError` to this list in the method implementation
+    in order for it to get picked up by the `KernelValidation` run.
+    """
+
+
+@dataclass
+class ValidationAnalysis(ForwardExtra[ValidationFrame, ErrorType], ABC):
+    """Analysis pass that indicates errors in the IR according to the respective method tables.
+
+    If you need to implement validation for a dialect shared by many groups (for example, if you need to ascertain if statements have a specific form)
+    you'll need to inherit from this class.
+    """
+
+    lattice = ErrorType
+
+    def eval_fallback(self, frame: ValidationFrame, node: ir.Statement):
+        # NOTE: default to no errors
+        return tuple(self.lattice.top() for _ in node.results)
+
+    def initialize_frame(
+        self, node: ir.Statement, *, has_parent_access: bool = False
+    ) -> ValidationFrame:
+        return ValidationFrame(node, has_parent_access=has_parent_access)
+
+    def method_self(self, method: ir.Method) -> ErrorType:
+        return self.lattice.top()

bloqade/validation/analysis/lattice.py

@@ -0,0 +1,58 @@
+from typing import final
+from dataclasses import dataclass
+
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    IsSubsetEqMixin,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+
+
+@dataclass
+class ErrorType(
+    SimpleJoinMixin["ErrorType"],
+    SimpleMeetMixin["ErrorType"],
+    IsSubsetEqMixin["ErrorType"],
+    BoundedLattice["ErrorType"],
+):
+
+    @classmethod
+    def bottom(cls) -> "ErrorType":
+        return InvalidErrorType()
+
+    @classmethod
+    def top(cls) -> "ErrorType":
+        return NoError()
+
+
+@final
+@dataclass
+class InvalidErrorType(ErrorType, metaclass=SingletonMeta):
+    """Bottom to represent when we encounter an error running the analysis.
+
+    When this is encountered, it means there might be an error, but we were unable to tell.
+    """
+
+    pass
+
+
+@final
+@dataclass
+class Error(ErrorType):
+    """Indicates an error in the IR."""
+
+    message: str = ""
+    """Optional error message to show in the IR.
+
+    NOTE: this is just to show a message when printing the IR. Actual errors
+    are collected by appending ir.ValidationError to the frame in the method
+    implementation.
+    """
+
+
+@final
+@dataclass
+class NoError(ErrorType, metaclass=SingletonMeta):
+    pass

bloqade/validation/kernel_validation.py

@@ -0,0 +1,77 @@
+import sys
+from dataclasses import dataclass
+
+from kirin import ir, exception
+from rich.console import Console
+
+from .analysis import ValidationAnalysis
+
+
+class ValidationErrorGroup(BaseException):
+    def __init__(self, *args: object, errors=[]) -> None:
+        super().__init__(*args)
+        self.errors = errors
+
+
+# TODO: this overrides kirin's exception handler and should be upstreamed
+def exception_handler(exc_type, exc_value, exc_tb):
+    if issubclass(exc_type, ValidationErrorGroup):
+        console = Console(force_terminal=True)
+        for i, err in enumerate(exc_value.errors):
+            with console.capture() as capture:
+                console.print(f"==== Error {i} ====")
+                console.print(f"[bold red]{type(err).__name__}:[/bold red]", end="")
+            print(capture.get(), *err.args, file=sys.stderr)
+            if err.source:
+                print("Source Traceback:", file=sys.stderr)
+                print(err.hint(), file=sys.stderr, end="")
+        console.print("=" * 40)
+        console.print(
+            "[bold red]Kernel validation failed:[/bold red] There were multiple errors encountered during validation, see above"
+        )
+        return
+
+    return exception.exception_handler(exc_type, exc_value, exc_tb)
+
+
+sys.excepthook = exception_handler
+
+
+@dataclass
+class KernelValidation:
+    """Validate a kernel according to a `ValidationAnalysis`.
+
+    This is a simple wrapper around the analysis that runs the analysis, checks
+    the `ValidationFrame` for errors and throws them if there are any.
+    """
+
+    validation_analysis_cls: type[ValidationAnalysis]
+    """The analysis that you want to run in order to validate the kernel."""
+
+    def run(self, mt: ir.Method, no_raise: bool = True) -> None:
+        """Run the kernel validation analysis and raise any errors found.
+
+        Args:
+            mt (ir.Method): The method to validate
+            no_raise (bool): Whether or not to raise errors when running the analysis.
+                This is only to make sure the analysis works. Errors found during
+                the analysis will be raised regardless of this setting. Defaults to `True`.
+
+        """
+
+        validation_analysis = self.validation_analysis_cls(mt.dialects)
+
+        if no_raise:
+            validation_frame, _ = validation_analysis.run_no_raise(mt)
+        else:
+            validation_frame, _ = validation_analysis.run(mt)
+
+        errors = validation_frame.errors
+
+        if len(errors) == 0:
+            # Valid program
+            return
+        elif len(errors) == 1:
+            raise errors[0]
+        else:
+            raise ValidationErrorGroup(errors=errors)

{bloqade_circuit-0.6.2.dist-info → bloqade_circuit-0.9.1.dist-info}/METADATA

@@ -1,16 +1,15 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.6.2
+Version: 0.9.1
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: kirin-toolchain~=0.17.0
+Requires-Dist: kirin-toolchain~=0.21.0
 Requires-Dist: numpy>=1.22.0
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: pydantic<2.11.0,>=1.3.0
-Requires-Dist: pyqrack-cpu<1.41,>=1.38.2; sys_platform != 'darwin'
-Requires-Dist: pyqrack<1.41,>=1.38.2; sys_platform == 'darwin'
+Requires-Dist: pyqrack-cpu~=1.69.1
 Requires-Dist: rich>=13.9.4
 Requires-Dist: scipy>=1.13.1
 Provides-Extra: cirq
@@ -18,9 +17,9 @@ Requires-Dist: cirq-core>=1.4.1; extra == 'cirq'
 Requires-Dist: cirq-core[contrib]>=1.4.1; extra == 'cirq'
 Requires-Dist: qpsolvers[clarabel]>=4.7.0; extra == 'cirq'
 Provides-Extra: pyqrack-cuda
-Requires-Dist: pyqrack-cuda>=1.38.2; extra == 'pyqrack-cuda'
+Requires-Dist: pyqrack-cuda~=1.69.1; extra == 'pyqrack-cuda'
 Provides-Extra: pyqrack-opencl
-Requires-Dist: pyqrack>=1.38.2; (sys_platform != 'darwin') and extra == 'pyqrack-opencl'
+Requires-Dist: pyqrack~=1.69.1; extra == 'pyqrack-opencl'
 Provides-Extra: qasm2
 Requires-Dist: lark>=1.2.2; extra == 'qasm2'
 Provides-Extra: qbraid