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

bloqade/squin/qubit.py

@@ -1,184 +0,0 @@
-"""qubit dialect for squin language.
-
-This dialect defines the operations that can be performed on qubits.
-
-Depends on:
-- `bloqade.squin.op`: provides the `OpType` type and semantics for operators applied to qubits.
-- `kirin.dialects.ilist`: provides the `ilist.IListType` type for lists of qubits.
-"""
-
-from typing import Any, overload
-
-from kirin import ir, types, lowering
-from kirin.decl import info, statement
-from kirin.dialects import ilist
-from kirin.lowering import wraps
-
-from bloqade.types import Qubit, QubitType
-from bloqade.squin.op.types import Op, OpType
-
-from .types import MeasurementResult, MeasurementResultType
-from .lowering import ApplyAnyCallLowering
-
-dialect = ir.Dialect("squin.qubit")
-
-
-@statement(dialect=dialect)
-class New(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    n_qubits: ir.SSAValue = info.argument(types.Int)
-    result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
-
-
-@statement(dialect=dialect)
-class Apply(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    operator: ir.SSAValue = info.argument(OpType)
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-
-
-@statement(dialect=dialect)
-class ApplyAny(ir.Statement):
-    # NOTE: custom lowering to deal with vararg calls
-    traits = frozenset({ApplyAnyCallLowering()})
-    operator: ir.SSAValue = info.argument(OpType)
-    qubits: tuple[ir.SSAValue, ...] = info.argument()
-
-
-@statement(dialect=dialect)
-class Broadcast(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    operator: ir.SSAValue = info.argument(OpType)
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-
-
-@statement(dialect=dialect)
-class MeasureAny(ir.Statement):
-    name = "measure"
-
-    traits = frozenset({lowering.FromPythonCall()})
-    input: ir.SSAValue = info.argument(types.Any)
-    result: ir.ResultValue = info.result(types.Any)
-
-
-@statement(dialect=dialect)
-class MeasureQubit(ir.Statement):
-    name = "measure.qubit"
-
-    traits = frozenset({lowering.FromPythonCall()})
-    qubit: ir.SSAValue = info.argument(QubitType)
-    result: ir.ResultValue = info.result(MeasurementResultType)
-
-
-@statement(dialect=dialect)
-class MeasureQubitList(ir.Statement):
-    name = "measure.qubit.list"
-
-    traits = frozenset({lowering.FromPythonCall()})
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-    result: ir.ResultValue = info.result(ilist.IListType[MeasurementResultType])
-
-
-# NOTE: no dependent types in Python, so we have to mark it Any...
-@wraps(New)
-def new(n_qubits: int) -> ilist.IList[Qubit, Any]:
-    """Create 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.
-    """
-    ...
-
-
-@overload
-def apply(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
-    """Apply an operator to a list of qubits.
-
-    Note, that when considering atom loss, lost qubits will be skipped.
-
-    Args:
-        operator: The operator to apply.
-        qubits: The list of qubits to apply the operator to. The size of the list
-            must be inferable and match the number of qubits expected by the operator.
-
-    Returns:
-        None
-    """
-    ...
-
-
-@overload
-def apply(operator: Op, *qubits: Qubit) -> None:
-    """Apply and operator to any number of qubits.
-
-    Note, that when considering atom loss, lost qubits will be skipped.
-
-    Args:
-        operator: The operator to apply.
-        *qubits: The qubits to apply the operator to. The number of qubits must
-            match the size of the operator.
-
-    Returns:
-        None
-    """
-    ...
-
-
-@wraps(ApplyAny)
-def apply(operator: Op, *qubits) -> None: ...
-
-
-@overload
-def measure(input: Qubit) -> MeasurementResult: ...
-@overload
-def measure(
-    input: ilist.IList[Qubit, Any] | list[Qubit],
-) -> ilist.IList[MeasurementResult, Any]: ...
-
-
-@wraps(MeasureAny)
-def measure(input: Any) -> Any:
-    """Measure a qubit or qubits in the list.
-
-    Args:
-        input: A qubit or a list of qubits to measure.
-
-    Returns:
-        bool | list[bool]: The result of the measurement. If a single qubit is measured,
-            a single boolean is returned. If a list of qubits is measured, a list of booleans
-            is returned.
-    """
-    ...
-
-
-@wraps(Broadcast)
-def broadcast(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
-    """Broadcast and apply an operator to a list of qubits. For example, an operator
-    that expects 2 qubits can be applied to a list of 2n qubits, where n is an integer > 0.
-
-    For controlled operators, the list of qubits is interpreted as sets of (controls, targets).
-    For example
-
-    ```
-    apply(CX, [q0, q1, q2, q3])
-    ```
-
-    is equivalent to
-
-    ```
-    apply(CX, [q0, q1])
-    apply(CX, [q2, q3])
-    ```
-
-    Args:
-        operator: The operator to broadcast and apply.
-        qubits: The list of qubits to broadcast and apply the operator to. The size of the list
-            must be inferable and match the number of qubits expected by the operator.
-
-    Returns:
-        None
-    """
-    ...

bloqade/squin/rewrite/canonicalize.py

@@ -1,60 +0,0 @@
-from typing import cast
-
-from kirin import ir
-from kirin.rewrite import abc
-from kirin.dialects import cf
-
-from .. import wire
-
-
-class CanonicalizeWired(abc.RewriteRule):
-    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
-
-        if (
-            not isinstance(node, wire.Wired)
-            or len(node.qubits) != 0
-            or (parent_region := node.parent_region) is None
-        ):
-            return abc.RewriteResult()
-
-        parent_block = cast(ir.Block, node.parent_block)
-
-        # the body doesn't contain any quantum operations so we can safely inline the
-        # body into the parent block
-
-        # move all statements after `node` in the current block into another block
-        after_block = ir.Block()
-
-        stmt = node.next_stmt
-        while stmt is not None:
-            stmt.detach()
-            after_block.stmts.append(stmt)
-            stmt = node.next_stmt
-
-        # remap all results of the node to the arguments of the after_block
-        for result in node.results:
-            arg = after_block.args.append_from(result.type, result.name)
-            result.replace_by(arg)
-
-        parent_block_idx = parent_region._block_idx[parent_block]
-        # insert goto of parent block to the body block of the node.
-        parent_region.blocks.insert(parent_block_idx + 1, after_block)
-        # insert all blocks of the body of the node after the parent region
-        # making sure to convert any yield statements to jump statements to the after_block
-        parent_block.stmts.append(
-            cf.Branch(
-                arguments=(),
-                successor=node.body.blocks[0],
-            )
-        )
-        for block in reversed(node.body.blocks):
-            block.detach()
-            if isinstance((yield_stmt := block.last_stmt), wire.Yield):
-                yield_stmt.replace_by(
-                    cf.Branch(yield_stmt.values, successor=after_block)
-                )
-
-            parent_region.blocks.insert(parent_block_idx + 1, block)
-
-        node.delete()
-        return abc.RewriteResult(has_done_something=True)

bloqade/squin/rewrite/desugar.py

@@ -1,124 +0,0 @@
-from kirin import ir, types
-from kirin.dialects import py, ilist
-from kirin.rewrite.abc import RewriteRule, RewriteResult
-
-from bloqade.squin.qubit import (
-    Apply,
-    ApplyAny,
-    QubitType,
-    MeasureAny,
-    MeasureQubit,
-    MeasureQubitList,
-)
-
-
-class MeasureDesugarRule(RewriteRule):
-    """
-    Desugar measure operations in the circuit.
-    """
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-
-        if not isinstance(node, MeasureAny):
-            return RewriteResult()
-
-        if node.input.type.is_subseteq(QubitType):
-            node.replace_by(
-                MeasureQubit(
-                    qubit=node.input,
-                )
-            )
-            return RewriteResult(has_done_something=True)
-        elif node.input.type.is_subseteq(ilist.IListType[QubitType, types.Any]):
-            node.replace_by(
-                MeasureQubitList(
-                    qubits=node.input,
-                )
-            )
-            return RewriteResult(has_done_something=True)
-
-        return RewriteResult()
-
-
-class ApplyDesugarRule(RewriteRule):
-    """
-    Desugar apply operators in the kernel.
-    """
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-
-        if not isinstance(node, ApplyAny):
-            return RewriteResult()
-
-        op = node.operator
-        qubits = node.qubits
-
-        if len(qubits) > 1 and all(q.type.is_subseteq(QubitType) for q in qubits):
-            (qubits_ilist_stmt := ilist.New(qubits)).insert_before(
-                node
-            )  # qubits is just a tuple of SSAValues
-            qubits_ilist = qubits_ilist_stmt.result
-
-        elif len(qubits) == 1 and qubits[0].type.is_subseteq(QubitType):
-            (qubits_ilist_stmt := ilist.New(qubits)).insert_before(node)
-            qubits_ilist = qubits_ilist_stmt.result
-
-        elif len(qubits) == 1 and qubits[0].type.is_subseteq(
-            ilist.IListType[QubitType, types.Any]
-        ):
-            qubits_ilist = qubits[0]
-
-        elif len(qubits) == 1:
-            # TODO: remove this elif clause once we're at kirin v0.18
-            # NOTE: this is a temporary workaround for kirin#408
-            # currently type inference fails here in for loops since the loop var
-            # is an IList for some reason
-
-            if not isinstance(qubits[0], ir.ResultValue):
-                return RewriteResult()
-
-            is_ilist = isinstance(qbit_stmt := qubits[0].stmt, ilist.New)
-
-            if is_ilist:
-
-                if not all(
-                    isinstance(qbit_getindex_result, ir.ResultValue)
-                    for qbit_getindex_result in qbit_stmt.values
-                ):
-                    return RewriteResult()
-
-                # Get the parent statement that the qubit came from
-                # (should be a GetItem instance, see logic below)
-                qbit_getindices = [
-                    qbit_getindex_result.stmt
-                    for qbit_getindex_result in qbit_stmt.values
-                ]
-            else:
-                qbit_getindices = [qubit.stmt for qubit in qubits]
-
-            if any(
-                not isinstance(qbit_getindex, py.indexing.GetItem)
-                for qbit_getindex in qbit_getindices
-            ):
-                return RewriteResult()
-
-            # The GetItem should have been applied on something that returns an IList of Qubits
-            if any(
-                not qbit_getindex.obj.type.is_subseteq(
-                    ilist.IListType[QubitType, types.Any]
-                )
-                for qbit_getindex in qbit_getindices
-            ):
-                return RewriteResult()
-
-            if is_ilist:
-                qubits_ilist = qbit_stmt.result
-            else:
-                (qubits_ilist_stmt := ilist.New(values=[qubits[0]])).insert_before(node)
-                qubits_ilist = qubits_ilist_stmt.result
-        else:
-            return RewriteResult()
-
-        stmt = Apply(operator=op, qubits=qubits_ilist)
-        node.replace_by(stmt)
-        return RewriteResult(has_done_something=True)

bloqade/squin/types.py

@@ -1,8 +0,0 @@
-from kirin import types
-
-
-class MeasurementResult:
-    pass
-
-
-MeasurementResultType = types.PyClass(MeasurementResult)

bloqade/squin/wire.py

@@ -1,201 +0,0 @@
-"""A NVIDIA QUAKE-like wire dialect.
-
-This dialect is expected to be used in combination with the operator dialect
-as an intermediate representation for analysis and optimization of quantum
-circuits. Thus we do not define wrapping functions for the statements in this
-dialect.
-"""
-
-from kirin import ir, types, lowering, exception
-from kirin.decl import info, statement
-from kirin.dialects import func
-from kirin.lowering import wraps
-from kirin.ir.attrs.types import TypeAttribute
-
-from bloqade.types import Qubit, QubitType
-
-from .types import MeasurementResultType
-from .op.types import Op, OpType
-
-# from kirin.lowering import wraps
-
-# from .op.types import Op, OpType
-
-dialect = ir.Dialect("squin.wire")
-
-
-class WireTerminator(ir.StmtTrait):
-    pass
-
-
-class Wire:
-    pass
-
-
-WireType = types.PyClass(Wire)
-
-
-# no return value for `wrap`
-@statement(dialect=dialect)
-class Wrap(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
-    wire: ir.SSAValue = info.argument(WireType)
-    qubit: ir.SSAValue = info.argument(QubitType)
-
-
-# "Unwrap the quantum references to expose wires" -> From Quake Dialect documentation
-# Unwrap(Qubit) -> Wire
-@statement(dialect=dialect)
-class Unwrap(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
-    qubit: ir.SSAValue = info.argument(QubitType)
-    result: ir.ResultValue = info.result(WireType)
-
-
-@statement(dialect=dialect)
-class Wired(ir.Statement):
-    traits = frozenset()
-
-    qubits: tuple[ir.SSAValue, ...] = info.argument(QubitType)
-    memory_zone: str = info.attribute()
-    body: ir.Region = info.region(multi=True)
-
-    def __init__(
-        self,
-        body: ir.Region,
-        *qubits: ir.SSAValue,
-        memory_zone: str,
-        result_types: tuple[TypeAttribute, ...] | None = None,
-    ):
-        if result_types is None:
-            for block in body.blocks:
-                if isinstance(block.last_stmt, Yield):
-                    result_types = tuple(arg.type for arg in block.last_stmt.values)
-                    break
-
-        if result_types is None:
-            result_types = ()
-
-        super().__init__(
-            args=qubits,
-            args_slice={
-                "qubits": slice(0, None),
-            },
-            regions=[body],
-            attributes={
-                "memory_zone": ir.PyAttr(memory_zone)
-            },  # body of the wired statement
-            result_types=result_types,
-        )
-
-    def check(self):
-        entry_block = self.body.blocks[0]
-
-        if len(entry_block.args) != len(self.qubits):
-            raise exception.StaticCheckError(
-                f"Expected {len(self.qubits)} arguments, got {len(entry_block.args)}."
-            )
-        for arg in entry_block.args:
-            if not arg.type.is_subseteq(WireType):
-                raise exception.StaticCheckError(
-                    f"Expected argument of type {WireType}, got {arg.type}."
-                )
-        for block in self.body.blocks:
-            last_stmt = block.last_stmt
-            if isinstance(last_stmt, func.Return):
-                raise exception.StaticCheckError(
-                    "Return statements are not allowed in the body of a Wired statement."
-                )
-            elif isinstance(last_stmt, Yield) and len(last_stmt.values) != len(
-                self.results
-            ):
-                raise exception.StaticCheckError(
-                    f"Expected {len(self.results)} return values, got {len(last_stmt.values)}."
-                )
-
-
-@statement(dialect=dialect)
-class Yield(ir.Statement):
-    traits = frozenset({})
-    values: tuple[ir.SSAValue, ...] = info.argument(WireType)
-
-    def __init__(self, *args: ir.SSAValue):
-        super().__init__(
-            args=args,
-            args_slice={
-                "values": slice(0, None),
-            },
-        )
-
-
-# In Quake, you put a wire in and get a wire out when you "apply" an operator
-# In this case though we just need to indicate that an operator is applied to list[wires]
-@statement(dialect=dialect)
-class Apply(ir.Statement):  # apply(op, w1, w2, ...)
-    traits = frozenset({lowering.FromPythonCall()})
-    operator: ir.SSAValue = info.argument(OpType)
-    inputs: tuple[ir.SSAValue, ...] = info.argument(WireType)
-
-    def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
-        result_types = tuple(WireType for _ in args)
-        super().__init__(
-            args=(operator,) + args,
-            result_types=result_types,  # result types of the Apply statement, should all be WireTypes
-            args_slice={
-                "operator": 0,
-                "inputs": slice(1, None),
-            },  # pretty printing + syntax sugar
-        )  # custom lowering required for wrapper to work here
-
-
-# Carry over from Qubit dialect
-@statement(dialect=dialect)
-class Broadcast(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
-    operator: ir.SSAValue = info.argument(OpType)
-    inputs: tuple[ir.SSAValue, ...] = info.argument(WireType)
-
-    def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
-        result_types = tuple(WireType for _ in args)
-        super().__init__(
-            args=(operator,) + args,
-            result_types=result_types,
-            args_slice={
-                "operator": 0,
-                "inputs": slice(1, None),
-            },  # pretty printing + syntax sugar
-        )  # custom lowering required for wrapper to work here
-
-
-@statement(dialect=dialect)
-class RegionMeasure(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
-    wire: ir.SSAValue = info.argument(WireType)
-    result: ir.ResultValue = info.result(MeasurementResultType)
-
-
-# NOTE: measurement cannot be pure because they will collapse the state
-#       of the qubit. The state is a hidden state that is not visible to
-#      the user in the wire dialect.
-@statement(dialect=dialect)
-class Measure(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
-    wire: ir.SSAValue = info.argument(WireType)
-    qubit: ir.SSAValue = info.argument(QubitType)
-    result: ir.ResultValue = info.result(MeasurementResultType)
-
-
-@statement(dialect=dialect)
-class LossResolvingMeasure(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    input_wire: ir.SSAValue = info.argument(WireType)
-    result: ir.ResultValue = info.result(MeasurementResultType)
-    out_wire: ir.ResultValue = info.result(WireType)
-
-
-@wraps(Unwrap)
-def unwrap(qubit: Qubit) -> Wire: ...
-
-
-@wraps(Apply)
-def apply(op: Op, w: Wire) -> Wire: ...

bloqade/stim/rewrite/wire_identity_elimination.py

@@ -1,24 +0,0 @@
-from kirin import ir
-from kirin.rewrite.abc import RewriteRule, RewriteResult
-
-from bloqade.squin import wire
-
-
-class SquinWireIdentityElimination(RewriteRule):
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        """
-        Handle the case where an unwrap feeds a wire directly into a wrap,
-        equivalent to nothing happening/identity operation
-
-        w = unwrap(qubit)
-        wrap(qubit, w)
-        """
-        if isinstance(node, wire.Wrap):
-            wire_origin_stmt = node.wire.owner
-            if isinstance(wire_origin_stmt, wire.Unwrap):
-                node.delete()  # get rid of wrap
-                wire_origin_stmt.delete()  # get rid of the unwrap
-                return RewriteResult(has_done_something=True)
-
-        return RewriteResult()

bloqade/stim/rewrite/wire_to_stim.py

@@ -1,57 +0,0 @@
-from kirin import ir
-from kirin.rewrite.abc import RewriteRule, RewriteResult
-
-from bloqade.squin import op, wire, noise
-from bloqade.stim.rewrite.util import (
-    SQUIN_STIM_OP_MAPPING,
-    rewrite_Control,
-    rewrite_QubitLoss,
-    insert_qubit_idx_from_wire_ssa,
-)
-
-
-class SquinWireToStim(RewriteRule):
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        match node:
-            case wire.Apply() | wire.Broadcast():
-                return self.rewrite_Apply_and_Broadcast(node)
-            case _:
-                return RewriteResult()
-
-    def rewrite_Apply_and_Broadcast(
-        self, stmt: wire.Apply | wire.Broadcast
-    ) -> RewriteResult:
-
-        # this is an SSAValue, need it to be the actual operator
-        applied_op = stmt.operator.owner
-
-        if isinstance(applied_op, noise.stmts.QubitLoss):
-            return rewrite_QubitLoss(stmt)
-
-        assert isinstance(applied_op, op.stmts.Operator)
-
-        if isinstance(applied_op, op.stmts.Control):
-            return rewrite_Control(stmt)
-
-        stim_1q_op = SQUIN_STIM_OP_MAPPING.get(type(applied_op))
-        if stim_1q_op is None:
-            return RewriteResult()
-
-        qubit_idx_ssas = insert_qubit_idx_from_wire_ssa(
-            wire_ssas=stmt.inputs, stmt_to_insert_before=stmt
-        )
-        if qubit_idx_ssas is None:
-            return RewriteResult()
-
-        stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
-
-        # Get the wires from the inputs of Apply or Broadcast,
-        # then put those as the result of the current stmt
-        # before replacing it entirely
-        for input_wire, output_wire in zip(stmt.inputs, stmt.results):
-            output_wire.replace_by(input_wire)
-
-        stmt.replace_by(stim_1q_stmt)
-
-        return RewriteResult(has_done_something=True)