bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/rewrite/parallel_to_uop.py

@@ -0,0 +1,83 @@
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite import abc, result
+
+from bloqade.analysis import address
+from bloqade.qasm2.dialects import uop, parallel
+
+
+@dataclass
+class ParallelToUOpRule(abc.RewriteRule):
+    id_map: Dict[int, ir.SSAValue]
+    address_analysis: Dict[ir.SSAValue, address.Address]
+
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        if type(node) in parallel.dialect.stmts:
+            return getattr(self, f"rewrite_{node.name}")(node)
+
+        return result.RewriteResult()
+
+    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):
+            return None
+
+        ids = []
+        for ele in addr.data:
+            if not isinstance(ele, address.AddressQubit):
+                return None
+
+            ids.append(ele.data)
+
+        return [self.id_map[ele] for ele in ids]
+
+    def rewrite_cz(self, node: ir.Statement):
+        assert isinstance(node, parallel.CZ)
+
+        ctrls = self.get_qubit_ssa(node.ctrls)
+        qargs = self.get_qubit_ssa(node.qargs)
+
+        if ctrls is None or qargs is None:
+            return result.RewriteResult()
+
+        for ctrl, qarg in zip(ctrls, qargs):
+            new_node = uop.CZ(ctrl, qarg)
+            new_node.insert_before(node)
+
+        node.delete()
+
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_u(self, node: ir.Statement):
+        assert isinstance(node, parallel.UGate)
+
+        qargs = self.get_qubit_ssa(node.qargs)
+
+        if qargs is None:
+            return result.RewriteResult()
+
+        for qarg in qargs:
+            new_node = uop.UGate(qarg, theta=node.theta, phi=node.phi, lam=node.lam)
+            new_node.insert_after(node)
+
+        node.delete()
+
+        return result.RewriteResult(has_done_something=True)
+
+    def rewrite_rz(self, node: ir.Statement):
+        assert isinstance(node, parallel.RZ)
+
+        qargs = self.get_qubit_ssa(node.qargs)
+
+        if qargs is None:
+            return result.RewriteResult()
+
+        for qarg in qargs:
+            new_node = uop.RZ(qarg, theta=node.theta)
+            new_node.insert_after(node)
+
+        node.delete()
+
+        return result.RewriteResult(has_done_something=True)

bloqade/qasm2/rewrite/register.py

@@ -0,0 +1,45 @@
+from kirin import ir
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.qasm2.dialects import core
+
+
+class RaiseRegisterRule(RewriteRule):
+    """This rule puts all registers at the top of the block.
+
+    This is required for the UOpToParallel rules to work correctly
+    to handle cases where a register is defined in between two statements
+    that can be parallelized.
+
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, core.QRegNew):
+            return RewriteResult()
+
+        if node.parent_block is None or node.parent_block.first_stmt is None:
+            return RewriteResult()
+
+        first_stmt = node.parent_block.first_stmt
+
+        n_qubits_ref = node.n_qubits
+
+        n_qubits = n_qubits_ref.owner
+        if isinstance(n_qubits, py.Constant):
+            # 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)
+            new_n_qubits_ref = new_n_qubits.result
+
+        elif isinstance(n_qubits, ir.BlockArgument):
+            # case where the n_qubits comes from a block argument
+            new_n_qubits_ref = n_qubits
+        else:
+            return RewriteResult()
+
+        new_qreg_stmt = core.QRegNew(n_qubits=new_n_qubits_ref)
+        new_qreg_stmt.insert_before(first_stmt)
+        node.result.replace_by(new_qreg_stmt.result)
+        node.delete()
+        return RewriteResult(has_done_something=True)

bloqade/qasm2/rewrite/uop_to_parallel.py

@@ -0,0 +1,395 @@
+import abc
+from typing import Dict, List, Tuple, Iterable
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.rewrite import abc as rewrite_abc
+from kirin.dialects import py, ilist
+from kirin.analysis.const import lattice
+
+from bloqade.analysis import address
+from bloqade.qasm2.dialects import uop, core, parallel
+from bloqade.squin.analysis.schedule import StmtDag
+
+
+class MergePolicyABC(abc.ABC):
+    @abc.abstractmethod
+    def __call__(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def can_merge(cls, stmt1: ir.Statement, stmt2: ir.Statement) -> bool:
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def merge_gates(
+        cls, gate_stmts: Iterable[ir.Statement]
+    ) -> List[List[ir.Statement]]:
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def from_analysis(
+        cls, dag: StmtDag, address_analysis: Dict[ir.SSAValue, address.Address]
+    ) -> "MergePolicyABC":
+        pass
+
+
+@dataclass
+class SimpleMergePolicy(MergePolicyABC):
+    """General merge policy for merging gates based on their type and arguments.
+
+    Base class to implement a merge policy for CZ, U and RZ gates, To completed the policy implement the
+    `merge_gates` class method. This will take an iterable of statements and return a list
+    of groups of statements that can be merged together. There are two mix-in classes
+    that can be used to implement the `merge_gates` method. The `GreedyMixin` will merge
+    gates together greedily, while the `OptimalMixIn` will merge gates together optimally.
+
+    """
+
+    address_analysis: Dict[ir.SSAValue, address.Address]
+    """Mapping from SSA values to their address analysis results. Needed for rewrites"""
+    merge_groups: List[List[ir.Statement]]
+    """List of groups of statements that can be merged together"""
+    group_numbers: Dict[ir.Statement, int]
+    """Mapping from statements to their group number"""
+    group_has_merged: Dict[int, bool] = field(default_factory=dict)
+    """Mapping from group number to whether the group has been merged"""
+
+    @staticmethod
+    def same_id_checker(ssa1: ir.SSAValue, ssa2: ir.SSAValue):
+        if ssa1 is ssa2:
+            return True
+        elif (hint1 := ssa1.hints.get("const")) and (hint2 := ssa2.hints.get("const")):
+            assert isinstance(hint1, lattice.Result) and isinstance(
+                hint2, lattice.Result
+            )
+            return hint1.is_equal(hint2)
+        else:
+            return False
+
+    @classmethod
+    def check_equiv_args(
+        cls,
+        args1: Iterable[ir.SSAValue],
+        args2: Iterable[ir.SSAValue],
+    ):
+        try:
+            return all(
+                cls.same_id_checker(ssa1, ssa2)
+                for ssa1, ssa2 in zip(args1, args2, strict=True)
+            )
+        except ValueError:
+            return False
+
+    @classmethod
+    def can_merge(cls, stmt1: ir.Statement, stmt2: ir.Statement) -> bool:
+        match stmt1, stmt2:
+            case (
+                (uop.UGate(), uop.UGate())
+                | (uop.RZ(), uop.RZ())
+                | (parallel.UGate(), parallel.UGate())
+                | (parallel.UGate(), uop.UGate())
+                | (uop.UGate(), parallel.UGate())
+                | (uop.UGate(), parallel.UGate())
+                | (uop.UGate(), parallel.UGate())
+                | (parallel.RZ(), parallel.RZ())
+                | (uop.RZ(), parallel.RZ())
+                | (parallel.RZ(), uop.RZ())
+            ):
+                return cls.check_equiv_args(stmt1.args[1:], stmt2.args[1:])
+            case (
+                (parallel.CZ(), parallel.CZ())
+                | (parallel.CZ(), uop.CZ())
+                | (uop.CZ(), parallel.CZ())
+                | (uop.CZ(), uop.CZ())
+                | (uop.Barrier(), uop.Barrier())
+            ):
+                return True
+
+            case _:
+                return False
+
+    @classmethod
+    def from_analysis(
+        cls,
+        dag: StmtDag,
+        address_analysis: Dict[ir.SSAValue, address.Address],
+    ):
+
+        merge_groups = []
+        group_numbers = {}
+
+        for group in dag.topological_groups():
+            gate_groups = cls.merge_gates(map(dag.stmts.__getitem__, group))
+            gate_groups_iter = (group for group in gate_groups if len(group) > 1)
+
+            for gate_group in gate_groups_iter:
+                group_number = len(merge_groups)
+                merge_groups.append(gate_group)
+                for stmt in gate_group:
+                    group_numbers[stmt] = group_number
+
+        for group in merge_groups:
+            group.sort(key=lambda stmt: dag.stmt_index[stmt])
+
+        return cls(
+            address_analysis=address_analysis,
+            merge_groups=merge_groups,
+            group_numbers=group_numbers,
+        )
+
+    def __call__(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+
+        if node not in self.group_numbers:
+            return rewrite_abc.RewriteResult()
+
+        group_number = self.group_numbers[node]
+        group = self.merge_groups[group_number]
+        if node is group[0]:
+            result = getattr(self, f"rewrite_group_{node.name}")(node, group)
+
+            self.group_has_merged[group_number] = result.has_done_something
+            return result
+
+        if self.group_has_merged[group_number]:
+            node.delete()
+
+        return rewrite_abc.RewriteResult(
+            has_done_something=self.group_has_merged[group_number]
+        )
+
+    def move_and_collect_qubit_list(
+        self, qargs: List[ir.SSAValue], node: ir.Statement
+    ) -> Tuple[ir.SSAValue, ...] | None:
+
+        qubits: List[ir.SSAValue] = []
+        # collect references to qubits
+        for qarg in qargs:
+            addr = self.address_analysis[qarg]
+
+            if isinstance(addr, address.AddressQubit):
+                qubits.append(qarg)
+
+            elif isinstance(addr, address.AddressTuple):
+                assert isinstance(qarg, ir.ResultValue)
+                assert isinstance(qarg.stmt, ilist.New)
+                qubits.extend(qarg.stmt.values)
+            else:
+                # give up if we cannot determine the address
+                return None
+
+        new_qubits = []
+
+        # the registers must be moved to the top of the block
+        # before this pass can be applied
+        for qubit_ref in qubits:
+            qubit = qubit_ref.owner
+            match qubit:
+                case ir.BlockArgument():  # do not need to move the qubit
+                    new_qubits.append(qubit)
+                case core.QRegGet(reg=reg, idx=ir.BlockArgument() as idx):
+                    new_qubit = core.QRegGet(reg=reg, idx=idx)
+                    new_qubit.insert_before(node)
+                    new_qubits.append(new_qubit.result)
+                case core.QRegGet(
+                    reg=reg, idx=ir.ResultValue(stmt=py.Constant() as idx)
+                ):
+                    (new_idx := idx.from_stmt(idx)).insert_before(node)
+                    (
+                        new_qubit := core.QRegGet(reg=reg, idx=new_idx.result)
+                    ).insert_before(node)
+                    new_qubits.append(new_qubit.result)
+                case _:
+                    return None
+
+        return tuple(new_qubits)
+
+    def rewrite_group_cz(self, node: ir.Statement, group: List[ir.Statement]):
+        ctrls = []
+        qargs = []
+
+        for stmt in group:
+            if isinstance(stmt, uop.CZ):
+                ctrls.append(stmt.ctrl)
+                qargs.append(stmt.qarg)
+            elif isinstance(stmt, parallel.CZ):
+                ctrls.append(stmt.ctrls)
+                qargs.append(stmt.qargs)
+            else:
+                return rewrite_abc.RewriteResult(has_done_something=False)
+
+        ctrls_values = self.move_and_collect_qubit_list(ctrls, node)
+        qargs_values = self.move_and_collect_qubit_list(qargs, node)
+
+        if ctrls_values is None or qargs_values is None:
+            # give up if we cannot determine the address or cannot move the qubits
+            return rewrite_abc.RewriteResult(has_done_something=False)
+
+        new_ctrls = ilist.New(values=ctrls_values)
+        new_qargs = ilist.New(values=qargs_values)
+        new_gate = parallel.CZ(ctrls=new_ctrls.result, qargs=new_qargs.result)
+
+        new_ctrls.insert_before(node)
+        new_qargs.insert_before(node)
+        new_gate.insert_before(node)
+
+        node.delete()
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+    def rewrite_group_U(self, node: ir.Statement, group: List[ir.Statement]):
+        return self.rewrite_group_u(node, group)
+
+    def rewrite_group_u(self, node: ir.Statement, group: List[ir.Statement]):
+        qargs = []
+
+        for stmt in group:
+            if isinstance(stmt, uop.UGate):
+                qargs.append(stmt.qarg)
+            elif isinstance(stmt, parallel.UGate):
+                qargs.append(stmt.qargs)
+            else:
+                return rewrite_abc.RewriteResult(has_done_something=False)
+
+        assert isinstance(node, (uop.UGate, parallel.UGate))
+        qargs_values = self.move_and_collect_qubit_list(qargs, node)
+
+        if qargs_values is None:
+            return rewrite_abc.RewriteResult(has_done_something=False)
+
+        new_qargs = ilist.New(values=qargs_values)
+        new_gate = parallel.UGate(
+            qargs=new_qargs.result,
+            theta=node.theta,
+            phi=node.phi,
+            lam=node.lam,
+        )
+        new_qargs.insert_before(node)
+        new_gate.insert_before(node)
+        node.delete()
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+    def rewrite_group_rz(self, node: ir.Statement, group: List[ir.Statement]):
+        qargs = []
+
+        for stmt in group:
+            if isinstance(stmt, uop.RZ):
+                qargs.append(stmt.qarg)
+            elif isinstance(stmt, parallel.RZ):
+                qargs.append(stmt.qargs)
+            else:
+                return rewrite_abc.RewriteResult(has_done_something=False)
+
+        assert isinstance(node, (uop.RZ, parallel.RZ))
+
+        qargs_values = self.move_and_collect_qubit_list(qargs, node)
+
+        if qargs_values is None:
+            return rewrite_abc.RewriteResult(has_done_something=False)
+
+        new_qargs = ilist.New(values=qargs_values)
+        new_gate = parallel.RZ(
+            qargs=new_qargs.result,
+            theta=node.theta,
+        )
+        new_qargs.insert_before(node)
+        new_gate.insert_before(node)
+        node.delete()
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+    def rewrite_group_barrier(self, node: uop.Barrier, group: List[uop.Barrier]):
+        qargs = []
+        for stmt in group:
+            qargs.extend(stmt.qargs)
+
+        qargs_values = self.move_and_collect_qubit_list(qargs, node)
+
+        if qargs_values is None:
+            return rewrite_abc.RewriteResult(has_done_something=False)
+
+        new_node = uop.Barrier(qargs=qargs_values)
+        new_node.insert_before(node)
+        node.delete()
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+
+class GreedyMixin(MergePolicyABC):
+    """Merge policy that greedily merges gates together.
+
+    The `merge_gates` method will merge policy will try greedily merge gates together.
+    This policy has a worst case complexity of O(n) where n is the
+    number of gates in the input iterable.
+    """
+
+    @classmethod
+    def merge_gates(
+        cls, gate_stmts: Iterable[ir.Statement]
+    ) -> List[List[ir.Statement]]:
+
+        iterable = iter(gate_stmts)
+        groups = [[next(iterable)]]
+
+        for stmt in gate_stmts:
+            if cls.can_merge(groups[-1][-1], stmt):
+                groups[-1].append(stmt)
+            else:
+                groups.append([stmt])
+
+        return groups
+
+
+class OptimalMixIn(MergePolicyABC):
+    """Merge policy that merges gates together optimally.
+
+    The `merge_gates` method will merge policy will try to merge every gate into every
+    group of gates, terminating when it finds a group that can be merged with the current
+    gate. This policy has a worst case complexity of O(n^2) where n is the number of gates
+    in the input iterable.
+
+    """
+
+    @classmethod
+    def merge_gates(
+        cls, gate_stmts: Iterable[ir.Statement]
+    ) -> List[List[ir.Statement]]:
+
+        groups = []
+        for stmt in gate_stmts:
+            found = False
+            for group in groups:
+                if cls.can_merge(group[-1], stmt):
+                    group.append(stmt)
+                    found = True
+                    break
+
+            if not found:
+                groups.append([stmt])
+
+        return groups
+
+
+@dataclass
+class SimpleGreedyMergePolicy(GreedyMixin, SimpleMergePolicy):
+    pass
+
+
+@dataclass
+class SimpleOptimalMergePolicy(OptimalMixIn, SimpleMergePolicy):
+    pass
+
+
+@dataclass
+class UOpToParallelRule(rewrite_abc.RewriteRule):
+    merge_rewriters: Dict[ir.Block | None, MergePolicyABC]
+
+    def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+        merge_rewriter = self.merge_rewriters.get(
+            node.parent_block, lambda _: rewrite_abc.RewriteResult()
+        )
+        return merge_rewriter(node)

bloqade/qasm2/types.py

@@ -0,0 +1,39 @@
+from kirin import types
+
+from bloqade.types import Qubit as Qubit, QubitType as QubitType
+
+
+class Bit:
+    """Runtime representation of a bit.
+
+    Note:
+        This is the base class of more specific bit types, such as
+        a reference to a piece of classical register in some quantum register
+        dialects.
+    """
+
+    pass
+
+
+class QReg:
+    """Runtime representation of a quantum register."""
+
+    def __getitem__(self, index) -> Qubit:
+        raise NotImplementedError("cannot call __getitem__ outside of a kernel")
+
+
+class CReg:
+    """Runtime representation of a classical register."""
+
+    def __getitem__(self, index) -> Bit:
+        raise NotImplementedError("cannot call __getitem__ outside of a kernel")
+
+
+BitType = types.PyClass(Bit)
+"""Kirin type for a classical bit."""
+
+QRegType = types.PyClass(QReg)
+"""Kirin type for a quantum register."""
+
+CRegType = types.PyClass(CReg)
+"""Kirin type for a classical register."""

bloqade/qbraid/__init__.py

@@ -0,0 +1,2 @@
+from .target import qBraid as qBraid
+from .lowering import Lowering as Lowering