bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qasm2/passes/noise.py

@@ -0,0 +1,61 @@
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    DeadCodeElimination,
+    CommonSubexpressionElimination,
+)
+from kirin.rewrite.result import RewriteResult
+
+from bloqade.noise import native
+from bloqade.analysis import address
+from bloqade.qasm2.rewrite.heuristic_noise import NoiseRewriteRule
+
+
+@dataclass
+class NoisePass(Pass):
+    """Apply a noise model to a quantum circuit.
+
+    NOTE: This pass is not guaranteed to be supported long-term in bloqade. We will be
+    moving towards a more general approach to noise modeling in the future.
+
+    """
+
+    noise_model: native.MoveNoiseModelABC = field(
+        default_factory=native.TwoRowZoneModel
+    )
+    gate_noise_params: native.GateNoiseParams = field(
+        default_factory=native.GateNoiseParams
+    )
+    address_analysis: address.AddressAnalysis = field(init=False)
+
+    def __post_init__(self):
+        self.address_analysis = address.AddressAnalysis(self.dialects)
+
+    def unsafe_run(self, mt: ir.Method):
+        result = RewriteResult()
+
+        frame, res = self.address_analysis.run_analysis(mt, no_raise=False)
+        result = (
+            Walk(
+                NoiseRewriteRule(
+                    address_analysis=frame.entries,
+                    noise_model=self.noise_model,
+                    gate_noise_params=self.gate_noise_params,
+                )
+            )
+            .rewrite(mt.code)
+            .join(result)
+        )
+        rule = Chain(
+            ConstantFold(),
+            DeadCodeElimination(),
+            CommonSubexpressionElimination(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+        return result

bloqade/qasm2/passes/parallel.py

@@ -0,0 +1,176 @@
+"""
+Passes for converting parallel gates into multiple single gates as well as
+converting multiple single gates to parallel gates.
+"""
+
+from typing import Type
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    WrapConst,
+    ConstantFold,
+    DeadCodeElimination,
+    CommonSubexpressionElimination,
+    result,
+)
+from kirin.analysis import const
+
+from bloqade.analysis import address
+from bloqade.qasm2.rewrite import (
+    MergePolicyABC,
+    ParallelToUOpRule,
+    RaiseRegisterRule,
+    UOpToParallelRule,
+    SimpleOptimalMergePolicy,
+)
+from bloqade.squin.analysis import schedule
+
+
+@dataclass
+class ParallelToUOp(Pass):
+    """Pass to convert parallel gates into single gates.
+
+    This pass rewrites any parallel gates from the `qasm2.parallel` dialect into multiple
+    single gates in the `qasm2.uop` dialect, bringing the program closer to
+    conforming to standard QASM2 syntax.
+
+    ## Usage Examples
+    ```
+    # Define kernel
+    @qasm2.extended
+    def main():
+        q = qasm2.qreg(4)
+
+        qasm2.parallel.cz(ctrls=[q[0], q[2]], qargs=[q[1], q[3]])
+
+    # Run rewrite
+    ParallelToUOp(main.dialects)(main)
+    ```
+
+    The `qasm2.parallel.cz` statement has been rewritten to individual gates:
+
+    ```
+    qasm2.uop.cz(ctrl=q[0], qarg=q[1])
+    qasm2.uop.cz(ctrl=q[2], qarg=q[3])
+    ```
+
+    """
+
+    def generate_rule(self, mt: ir.Method) -> ParallelToUOpRule:
+        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+
+        id_map = {}
+
+        # GOAL: Get the ssa value for the first reference of each qubit.
+        for ssa, addr in frame.entries.items():
+            if not isinstance(addr, address.AddressQubit):
+                # skip any stmts that are not qubits
+                continue
+
+            # get qubit id from analysis result
+            qubit_id = addr.data
+
+            # check if id has already been found
+            # if so, skip this ssa value
+            if qubit_id in id_map:
+                continue
+
+            id_map[qubit_id] = ssa
+
+        return ParallelToUOpRule(id_map=id_map, address_analysis=frame.entries)
+
+    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+        result = Walk(self.generate_rule(mt)).rewrite(mt.code)
+        rule = Chain(
+            ConstantFold(),
+            DeadCodeElimination(),
+            CommonSubexpressionElimination(),
+        )
+        return Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+
+@dataclass
+class UOpToParallel(Pass):
+    """Pass to convert single gates into parallel gates.
+
+    This pass looks for single gates from the `qasm2.uop` dialect that can be combined
+    into parallel gates from the `qasm2.parallel` dialect and performs a rewrite to do so.
+
+    ## Usage Examples
+    ```
+    # Define kernel
+    @qasm2.main
+    def test():
+        q = qasm2.qreg(4)
+
+        theta = 0.1
+        phi = 0.2
+        lam = 0.3
+
+        qasm2.u(q[1], theta, phi, lam)
+        qasm2.u(q[3], theta, phi, lam)
+        qasm2.cx(q[1], q[3])
+        qasm2.u(q[2], theta, phi, lam)
+        qasm2.u(q[0], theta, phi, lam)
+        qasm2.cx(q[0], q[2])
+
+    # Run rewrite
+    UOpToParallel(main.dialects)(main)
+    ```
+
+    The individual `qasm2.u` statements have now been combined
+    into a single `qasm2.parallel.u` statement.
+
+    ```
+    qasm2.parallel.u(qargs = [q[0], q[1], q[2], q[3]], theta, phi, lam)
+    qasm2.uop.CX(q[1], q[3])
+    qasm2.uop.CX(q[0], q[2])
+    ```
+
+    """
+
+    merge_policy_type: Type[MergePolicyABC] = SimpleOptimalMergePolicy
+    constprop: const.Propagate = field(init=False)
+
+    def __post_init__(self):
+        self.constprop = const.Propagate(self.dialects)
+
+    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+        result = Walk(RaiseRegisterRule()).rewrite(mt.code)
+
+        # do not run the parallelization because registers are not at the top
+        if not result.has_done_something:
+            return result
+
+        frame, _ = self.constprop.run_analysis(mt)
+        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
+
+        frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
+        dags = schedule.DagScheduleAnalysis(
+            mt.dialects, address_analysis=frame.entries
+        ).get_dags(mt)
+
+        result = (
+            Walk(
+                UOpToParallelRule(
+                    {
+                        block: self.merge_policy_type.from_analysis(dag, frame.entries)
+                        for block, dag in dags.items()
+                    }
+                )
+            )
+            .rewrite(mt.code)
+            .join(result)
+        )
+
+        rule = Chain(
+            ConstantFold(),
+            DeadCodeElimination(),
+            CommonSubexpressionElimination(),
+        )
+        return Fixpoint(Walk(rule)).rewrite(mt.code).join(result)

bloqade/qasm2/passes/py2qasm.py

@@ -0,0 +1,63 @@
+"""Rewrite py dialects into qasm dialects."""
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import Walk, Fixpoint
+from kirin.dialects import py, math
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+
+from bloqade.qasm2.dialects import core, expr
+
+
+class _Py2QASM(RewriteRule):
+    """Rewrite py dialects into qasm dialects."""
+
+    UNARY_OPS = {
+        py.USub: expr.Neg,
+        math.sin: expr.Sin,
+        math.cos: expr.Cos,
+        math.tan: expr.Tan,
+        math.exp: expr.Exp,
+        math.sqrt: expr.Sqrt,
+    }
+
+    BINARY_OPS = {
+        py.Add: expr.Add,
+        py.Sub: expr.Sub,
+        py.Mult: expr.Mul,
+        py.Div: expr.Div,
+        py.Pow: expr.Pow,
+    }
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if isinstance(node, py.Constant):
+            value = node.value.unwrap()
+            if isinstance(value, int):
+                node.replace_by(expr.ConstInt(value=value))
+                return RewriteResult(has_done_something=True)
+            elif isinstance(value, float):
+                node.replace_by(expr.ConstFloat(value=value))
+                return RewriteResult(has_done_something=True)
+        elif isinstance(node, py.BinOp):
+            if (pystmt := self.BINARY_OPS.get(type(node))) is not None:
+                node.replace_by(pystmt(node.lhs, node.rhs))
+                return RewriteResult(has_done_something=True)
+        elif isinstance(node, py.UnaryOp):
+            if (pystmt := self.UNARY_OPS.get(type(node))) is not None:
+                node.replace_by(pystmt(node.value))
+                return RewriteResult(has_done_something=True)
+        elif isinstance(node, py.cmp.Eq):
+            node.replace_by(core.CRegEq(node.lhs, node.rhs))
+            return RewriteResult(has_done_something=True)
+        elif isinstance(node, py.assign.Alias):
+            node.result.replace_by(node.value)
+            node.delete()
+            return RewriteResult(has_done_something=True)
+        return RewriteResult()
+
+
+class Py2QASM(Pass):
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        return Fixpoint(Walk(_Py2QASM())).rewrite(mt.code)

bloqade/qasm2/passes/qasm2py.py

@@ -0,0 +1,61 @@
+"""Rewrite qasm dialects into py dialects."""
+
+import math as pymath
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import Walk, Fixpoint
+from kirin.dialects import py, math
+from kirin.rewrite.abc import RewriteRule
+from kirin.rewrite.result import RewriteResult
+
+from bloqade.qasm2.dialects import core, expr
+
+
+class _QASM2Py(RewriteRule):
+    """Rewrite qasm dialects into py dialects."""
+
+    UNARY_OPS = {
+        expr.Neg: py.USub,
+        expr.Sin: math.stmts.sin,
+        expr.Cos: math.stmts.cos,
+        expr.Tan: math.stmts.tan,
+        expr.Exp: math.stmts.exp,
+        expr.Sqrt: math.stmts.sqrt,
+    }
+
+    BINARY_OPS = {
+        expr.Add: py.Add,
+        expr.Sub: py.Sub,
+        expr.Mul: py.Mult,
+        expr.Div: py.Div,
+        expr.Pow: py.Pow,
+    }
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if isinstance(node, (expr.ConstInt, expr.ConstFloat)):
+            node.replace_by(py.Constant(value=node.value))
+            return RewriteResult(has_done_something=True)
+        elif isinstance(node, expr.Neg):
+            node.replace_by(self.UNARY_OPS[type(node)](value=node.value))
+            return RewriteResult(has_done_something=True)
+        elif isinstance(node, (expr.Sin, expr.Cos, expr.Tan, expr.Exp, expr.Sqrt)):
+            node.replace_by(self.UNARY_OPS[type(node)](x=node.value))
+            return RewriteResult(has_done_something=True)
+        elif isinstance(node, (expr.Add, expr.Sub, expr.Mul, expr.Div, expr.Pow)):
+            node.replace_by(self.BINARY_OPS[type(node)](lhs=node.lhs, rhs=node.rhs))
+            return RewriteResult(has_done_something=True)
+        elif isinstance(node, core.CRegEq):
+            node.replace_by(py.cmp.Eq(node.lhs, node.rhs))
+            return RewriteResult(has_done_something=True)
+        elif isinstance(node, expr.ConstPI):
+            node.replace_by(py.Constant(value=pymath.pi))
+            return RewriteResult(has_done_something=True)
+        else:
+            return RewriteResult()
+
+
+class QASM2Py(Pass):
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        return Fixpoint(Walk(_QASM2Py())).rewrite(mt.code)

bloqade/qasm2/rewrite/__init__.py

@@ -0,0 +1,12 @@
+from .glob import (
+    GlobalToUOpRule as GlobalToUOpRule,
+    GlobalToParallelRule as GlobalToParallelRule,
+)
+from .register import RaiseRegisterRule as RaiseRegisterRule
+from .parallel_to_uop import ParallelToUOpRule as ParallelToUOpRule
+from .uop_to_parallel import (
+    MergePolicyABC as MergePolicyABC,
+    UOpToParallelRule as UOpToParallelRule,
+    SimpleGreedyMergePolicy as SimpleGreedyMergePolicy,
+    SimpleOptimalMergePolicy as SimpleOptimalMergePolicy,
+)

bloqade/qasm2/rewrite/desugar.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import abc, walk, result
+from kirin.dialects import py
+
+from bloqade.qasm2.dialects import core
+
+
+class IndexingDesugarRule(abc.RewriteRule):
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        if isinstance(node, py.indexing.GetItem):
+            if node.obj.type.is_subseteq(core.QRegType):
+                node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
+                return result.RewriteResult(has_done_something=True)
+            elif node.obj.type.is_subseteq(core.CRegType):
+                node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
+                return result.RewriteResult(has_done_something=True)
+
+        return result.RewriteResult()
+
+
+@dataclass
+class IndexingDesugarPass(Pass):
+    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+
+        return walk.Walk(IndexingDesugarRule()).rewrite(mt.code)

bloqade/qasm2/rewrite/glob.py

@@ -0,0 +1,103 @@
+from typing import Dict, List
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite import abc, result
+from kirin.dialects import py, ilist
+
+from bloqade import qasm2
+from bloqade.analysis import address
+from bloqade.qasm2.dialects import glob
+
+
+@dataclass
+class GlobalRewriteBase:
+    address_analysis: Dict[ir.SSAValue, address.Address]
+
+    def get_qubit_ssa(self, node: glob.UGate):
+        new_stmts: List[ir.Statement] = []
+        qubit_ssa: List[ir.SSAValue] = []
+        # can't rewrite if the registers are coming from a block argument
+        if not isinstance(node.registers, ir.ResultValue):
+            return new_stmts, None
+
+        if not isinstance(node.registers.owner, ilist.New):
+            return new_stmts, None
+
+        register_ssa_values = node.registers.owner.values
+
+        for register_ssa in register_ssa_values:
+            addr = self.address_analysis.get(register_ssa, address.Address.top())
+            if not isinstance(addr, address.AddressReg):
+                new_stmts.clear()
+                return new_stmts, None
+
+            for qubit in range(len(addr.data)):
+                new_stmts.append(idx_stmt := py.constant.Constant(value=qubit))
+                new_stmts.append(
+                    qubit_stmt := qasm2.core.QRegGet(
+                        reg=register_ssa, idx=idx_stmt.result
+                    )
+                )
+                qubit_ssa.append(qubit_stmt.result)
+
+        return new_stmts, qubit_ssa
+
+
+@dataclass
+class GlobalToParallelRule(abc.RewriteRule, GlobalRewriteBase):
+
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        if type(node) in glob.dialect.stmts:
+            return getattr(self, f"rewrite_{node.name}")(node)
+
+        return result.RewriteResult()
+
+    def rewrite_ugate(self, node: glob.UGate):
+
+        new_stmts, qubit_ssa = self.get_qubit_ssa(node)
+
+        if qubit_ssa is None:
+            return result.RewriteResult()
+
+        new_stmts.append(qargs := ilist.New(values=qubit_ssa))
+        new_stmts.append(
+            qasm2.dialects.parallel.UGate(
+                qargs=qargs.result, theta=node.theta, phi=node.phi, lam=node.lam
+            )
+        )
+
+        for stmt in new_stmts:
+            stmt.insert_before(node)
+
+        node.delete()
+
+        return result.RewriteResult(has_done_something=True)
+
+
+@dataclass
+class GlobalToUOpRule(abc.RewriteRule, GlobalRewriteBase):
+
+    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+        if type(node) in glob.dialect.stmts:
+            return getattr(self, f"rewrite_{node.name}")(node)
+
+        return result.RewriteResult()
+
+    def rewrite_ugate(self, node: glob.UGate):
+
+        new_stmts, qubit_ssa = self.get_qubit_ssa(node)
+
+        if qubit_ssa is None:
+            return result.RewriteResult()
+
+        for qarg in qubit_ssa:
+            new_stmts.append(
+                qasm2.uop.UGate(qarg=qarg, theta=node.theta, phi=node.phi, lam=node.lam)
+            )
+
+        for stmt in new_stmts:
+            stmt.insert_before(node)
+
+        node.delete()
+        return result.RewriteResult(has_done_something=True)