bloqade-circuit 0.2.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

bloqade/qasm2/parse/lowering.py

@@ -1,5 +1,3 @@
-import os
-import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -19,131 +17,39 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     hint_show_lineno: bool = field(default=True, kw_only=True)
     stacktrace: bool = field(default=True, kw_only=True)
 
-    def loads(
+    def run(
         self,
-        source: str,
-        kernel_name: str,
+        stmt: ast.Node,
         *,
-        returns: str | None = None,
+        source: str | None = None,
         globals: dict[str, Any] | None = None,
         file: str | None = None,
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
-    ) -> ir.Method:
-        from ..parse import loads
-
-        # TODO: add source info
-        stmt = loads(source)
+    ) -> ir.Region:
 
-        state = lowering.State(
-            self,
+        frame = self.get_frame(
+            stmt,
+            source=source,
+            globals=globals,
             file=file,
             lineno_offset=lineno_offset,
             col_offset=col_offset,
         )
-        with state.frame(
-            [stmt],
-            globals=globals,
-            finalize_next=False,
-        ) as frame:
-            try:
-                self.visit(state, stmt)
-                # append return statement with the return values
-                if returns is not None:
-                    return_value = frame.get(returns)
-                    if return_value is None:
-                        raise lowering.BuildError(f"Cannot find return value {returns}")
-                else:
-                    return_value = func.ConstantNone()
-                    frame.push(return_value)
-
-                return_node = frame.push(func.Return(value_or_stmt=return_value))
-
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-
-        code = func.Function(
-            sym_name=kernel_name,
-            signature=func.Signature((), return_node.value.type),
-            body=region,
-        )
 
-        return ir.Method(
-            mod=None,
-            py_func=None,
-            sym_name=kernel_name,
-            arg_names=[],
-            dialects=self.dialects,
-            code=code,
-        )
+        return frame.curr_region
 
-    def loadfile(
-        self,
-        file: str | pathlib.Path,
-        *,
-        kernel_name: str | None = None,
-        returns: str | None = None,
-        globals: dict[str, Any] | None = None,
-        lineno_offset: int = 0,
-        col_offset: int = 0,
-        compactify: bool = True,
-    ) -> ir.Method:
-        if isinstance(file, str):
-            file = pathlib.Path(*os.path.split(file))
-
-        if not file.is_file() or not file.name.endswith(".qasm"):
-            raise ValueError("File must be a .qasm file")
-
-        kernel_name = (
-            file.name.replace(".qasm", "") if kernel_name is None else kernel_name
-        )
-
-        with file.open("r") as f:
-            source = f.read()
-
-        return self.loads(
-            source,
-            kernel_name,
-            returns=returns,
-            globals=globals,
-            file=str(file),
-            lineno_offset=lineno_offset,
-            col_offset=col_offset,
-            compactify=compactify,
-        )
-
-    def run(
+    def get_frame(
         self,
         stmt: ast.Node,
-        *,
         source: str | None = None,
         globals: dict[str, Any] | None = None,
         file: str | None = None,
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
-    ) -> ir.Region:
+    ) -> lowering.Frame:
         # TODO: add source info
         state = lowering.State(
             self,
@@ -154,32 +60,16 @@ class QASM2(lowering.LoweringABC[ast.Node]):
         with state.frame(
             [stmt],
             globals=globals,
+            finalize_next=False,
         ) as frame:
-            try:
-                self.visit(state, stmt)
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-        return region
+            self.visit(state, stmt)
+
+            if compactify:
+                from kirin.rewrite import Walk, CFGCompactify
+
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
+
+            return frame
 
     def visit(self, state: lowering.State[ast.Node], node: ast.Node) -> lowering.Result:
         name = node.__class__.__name__

bloqade/qasm2/parse/qasm2.lark

@@ -9,7 +9,7 @@ version: INT "." INT
 // stmts
 include: "include" STRING ";"
 ifstmt: "if" "(" expr "==" expr ")" ifbody
-ifbody: qop | "{" qop* "}" // allow multiple qops
+ifbody: qop
 opaque: "opaque" IDENTIFIER ["(" [params] ")"] qubits ";"
 barrier: "barrier" qubits ";"
 qreg: "qreg" IDENTIFIER "[" INT "]" ";"

bloqade/qasm2/passes/__init__.py

@@ -3,3 +3,4 @@ from .noise import NoisePass as NoisePass
 from .py2qasm import Py2QASM as Py2QASM
 from .qasm2py import QASM2Py as QASM2Py
 from .parallel import UOpToParallel as UOpToParallel
+from .unroll_if import UnrollIfs as UnrollIfs

bloqade/qasm2/passes/fold.py

@@ -23,6 +23,8 @@ from kirin.rewrite.abc import RewriteResult
 
 from bloqade.qasm2.dialects import expr
 
+from .unroll_if import UnrollIfs
+
 
 @dataclass
 class QASM2Fold(Pass):
@@ -30,6 +32,7 @@ class QASM2Fold(Pass):
 
     constprop: const.Propagate = field(init=False)
     inline_gate_subroutine: bool = True
+    unroll_ifs: bool = True
 
     def __post_init__(self):
         self.constprop = const.Propagate(self.dialects)
@@ -61,6 +64,9 @@ class QASM2Fold(Pass):
             .join(result)
         )
 
+        if self.unroll_ifs:
+            UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
+
         # run typeinfer again after unroll etc. because we now insert
         # a lot of new nodes, which might have more precise types
         self.typeinfer.unsafe_run(mt)

bloqade/qasm2/passes/glob.py

@@ -58,13 +58,15 @@ class GlobalToUOP(Pass):
         rewriter = walk.Walk(self.generate_rule(mt))
         result = rewriter.rewrite(mt.code)
 
-        result = walk.Walk(dce.DeadCodeElimination()).rewrite(mt.code)
-        result = Fixpoint(walk.Walk(rule=cse.CommonSubexpressionElimination())).rewrite(
-            mt.code
+        result = walk.Walk(dce.DeadCodeElimination()).rewrite(mt.code).join(result)
+        result = (
+            Fixpoint(walk.Walk(rule=cse.CommonSubexpressionElimination()))
+            .rewrite(mt.code)
+            .join(result)
         )
 
         # do fold again to get proper hint for inserted const
-        result = Fold(mt.dialects)(mt)
+        result = Fold(mt.dialects)(mt).join(result)
         return result
 
 
@@ -110,10 +112,12 @@ class GlobalToParallel(Pass):
         rewriter = walk.Walk(self.generate_rule(mt))
         result = rewriter.rewrite(mt.code)
 
-        result = walk.Walk(dce.DeadCodeElimination()).rewrite(mt.code)
-        result = Fixpoint(walk.Walk(rule=cse.CommonSubexpressionElimination())).rewrite(
-            mt.code
+        result = walk.Walk(dce.DeadCodeElimination()).rewrite(mt.code).join(result)
+        result = (
+            Fixpoint(walk.Walk(rule=cse.CommonSubexpressionElimination()))
+            .rewrite(mt.code)
+            .join(result)
         )
         # do fold again to get proper hint
-        result = Fold(mt.dialects)(mt)
+        result = Fold(mt.dialects)(mt).join(result)
         return result

bloqade/qasm2/passes/noise.py

@@ -8,10 +8,10 @@ from kirin.rewrite import (
     DeadCodeElimination,
 )
 
-from bloqade.noise import native
+from bloqade.qasm2 import noise
 from bloqade.analysis import address
+from bloqade.qasm2.rewrite import NoiseRewriteRule
 from bloqade.qasm2.passes.lift_qubits import LiftQubits
-from bloqade.qasm2.rewrite.heuristic_noise import NoiseRewriteRule
 
 
 @dataclass
@@ -25,12 +25,9 @@ class NoisePass(Pass):
 
     ```
     from bloqade import qasm2
-    from bloqade.noise import native
-    from bloqade.qasm2.passes.noise import NoisePass
+    from bloqade.qasm2.passes import NoisePass
 
-    noise_main = qasm2.extended.add(native.dialect)
-
-    @noise_main
+    @qasm2.extended
     def main():
         q = qasm2.qreg(2)
         qasm2.h(q[0])
@@ -51,31 +48,45 @@ class NoisePass(Pass):
 
     """
 
-    noise_model: native.MoveNoiseModelABC = field(
-        default_factory=native.TwoRowZoneModel
-    )
-    gate_noise_params: native.GateNoiseParams = field(
-        default_factory=native.GateNoiseParams
-    )
+    noise_model: noise.MoveNoiseModelABC = field(default_factory=noise.TwoRowZoneModel)
     address_analysis: address.AddressAnalysis = field(init=False)
 
     def __post_init__(self):
         self.address_analysis = address.AddressAnalysis(self.dialects)
 
+    def get_qubit_values(self, mt: ir.Method):
+        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_values = {}
+        # Traverse statements in block order to fine the first SSA value for each qubit
+        for block in mt.callable_region.blocks:
+            for stmt in block.stmts:
+                if len(stmt.results) != 1:
+                    continue
+
+                addr = frame.entries.get(result := stmt.results[0])
+                if (
+                    isinstance(addr, address.AddressQubit)
+                    and (index := addr.data) not in qubit_ssa_values
+                ):
+                    qubit_ssa_values[index] = result
+
+        return qubit_ssa_values, frame.entries
+
     def unsafe_run(self, mt: ir.Method):
         result = LiftQubits(self.dialects).unsafe_run(mt)
-        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_value, address_analysis = self.get_qubit_values(mt)
         result = (
             Walk(
                 NoiseRewriteRule(
-                    address_analysis=frame.entries,
+                    qubit_ssa_value=qubit_ssa_value,
+                    address_analysis=address_analysis,
                     noise_model=self.noise_model,
-                    gate_noise_params=self.gate_noise_params,
                 ),
                 reverse=True,
             )
             .rewrite(mt.code)
             .join(result)
         )
+
         result = Fixpoint(Walk(DeadCodeElimination())).rewrite(mt.code).join(result)
         return result

bloqade/qasm2/passes/parallel.py

@@ -27,6 +27,7 @@ from bloqade.qasm2.rewrite import (
     RaiseRegisterRule,
     UOpToParallelRule,
     SimpleOptimalMergePolicy,
+    RydbergGateSetRewriteRule,
 )
 from bloqade.squin.analysis import schedule
 
@@ -135,6 +136,7 @@ class UOpToParallel(Pass):
     """
 
     merge_policy_type: Type[MergePolicyABC] = SimpleOptimalMergePolicy
+    rewrite_to_native_first: bool = False
     constprop: const.Propagate = field(init=False)
 
     def __post_init__(self):
@@ -147,6 +149,13 @@ class UOpToParallel(Pass):
         if not result.has_done_something:
             return result
 
+        if self.rewrite_to_native_first:
+            result = (
+                Fixpoint(Walk(RydbergGateSetRewriteRule(self.dialects)))
+                .rewrite(mt.code)
+                .join(result)
+            )
+
         frame, _ = self.constprop.run_analysis(mt)
         result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
 

bloqade/qasm2/passes/unroll_if.py

@@ -0,0 +1,25 @@
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    CommonSubexpressionElimination,
+)
+
+from ..rewrite.split_ifs import LiftThenBody, SplitIfStmts
+
+
+class UnrollIfs(Pass):
+    """This pass lifts statements that are not UOP out of the if body and then splits whatever is left into multiple if statements so you obtain valid QASM2"""
+
+    def unsafe_run(self, mt: ir.Method):
+        result = Walk(LiftThenBody()).rewrite(mt.code)
+        result = Walk(SplitIfStmts()).rewrite(mt.code).join(result)
+        result = (
+            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        return result

bloqade/qasm2/rewrite/__init__.py

@@ -3,6 +3,7 @@ from .glob import (
     GlobalToParallelRule as GlobalToParallelRule,
 )
 from .register import RaiseRegisterRule as RaiseRegisterRule
+from .native_gates import RydbergGateSetRewriteRule as RydbergGateSetRewriteRule
 from .parallel_to_uop import ParallelToUOpRule as ParallelToUOpRule
 from .uop_to_parallel import (
     MergePolicyABC as MergePolicyABC,
@@ -10,3 +11,5 @@ from .uop_to_parallel import (
     SimpleGreedyMergePolicy as SimpleGreedyMergePolicy,
     SimpleOptimalMergePolicy as SimpleOptimalMergePolicy,
 )
+from .noise.remove_noise import RemoveNoisePass as RemoveNoisePass
+from .noise.heuristic_noise import NoiseRewriteRule as NoiseRewriteRule

bloqade/qasm2/rewrite/desugar.py

@@ -5,16 +5,17 @@ from kirin.passes import Pass
 from kirin.rewrite import abc, walk
 from kirin.dialects import py
 
+from bloqade.qasm2 import types
 from bloqade.qasm2.dialects import core
 
 
 class IndexingDesugarRule(abc.RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, py.indexing.GetItem):
-            if node.obj.type.is_subseteq(core.QRegType):
+            if node.obj.type.is_subseteq(types.QRegType):
                 node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
-            elif node.obj.type.is_subseteq(core.CRegType):
+            elif node.obj.type.is_subseteq(types.CRegType):
                 node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
 

bloqade/qasm2/rewrite/native_gates.py

@@ -173,6 +173,53 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
             cirq.XPowGate(exponent=0.5).on(self.cached_qubits[0]), node
         )
 
+    def rewrite_ccx(self, node: uop.CCX) -> abc.RewriteResult:
+        # from https://algassert.com/quirk#circuit=%7B%22cols%22:%5B%5B%22QFT3%22%5D,%5B%22inputA3%22,1,1,%22+=A3%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22%E2%80%A2%22,%22X%22%5D,%5B1,1,1,%22%E2%80%A6%22,%22%E2%80%A6%22,%22%E2%80%A6%22%5D,%5B1,1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E-%C2%BC%22%5D,%5B1,1,1,%22%E2%80%A2%22,1,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E%C2%BC%22%5D,%5B1,1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E-%C2%BC%22%5D,%5B1,1,1,%22Z%5E%C2%BC%22,%22Z%5E%C2%BC%22%5D,%5B1,1,1,1,%22H%22%5D,%5B1,1,1,%22%E2%80%A2%22,1,%22Z%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,%22X%5E-%C2%BC%22,%22X%5E%C2%BC%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,%22H%22%5D%5D%7D
+
+        # x^(1/4)
+        lam1, theta1, phi1 = map(
+            self.const_float,
+            map(around, (1.5707963267948966, 0.7853981633974483, -1.5707963267948966)),
+        )
+        lam1.insert_before(node)
+        theta1.insert_before(node)
+        phi1.insert_before(node)
+
+        lam1 = lam1.result
+        theta1 = theta1.result
+        phi1 = phi1.result
+
+        # x^(-1/4)
+        lam2, theta2, phi2 = map(
+            self.const_float,
+            map(around, (4.71238898038469, 0.7853981633974483, 1.5707963267948966)),
+        )
+        lam2.insert_before(node)
+        theta2.insert_before(node)
+        phi2.insert_before(node)
+        lam2 = lam2.result
+        theta2 = theta2.result
+        phi2 = phi2.result
+
+        uop.CZ(ctrl=node.ctrl1, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta1, phi1, lam1).insert_before(node)
+        uop.CZ(ctrl=node.ctrl1, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.T(node.ctrl1).insert_before(node)
+        uop.T(node.ctrl2).insert_before(node)
+        uop.H(node.ctrl1).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.qarg).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.ctrl1).insert_before(node)
+        uop.UGate(node.ctrl1, theta2, phi2, lam2).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.ctrl1).insert_before(node)
+        uop.H(node.ctrl1).insert_before(node)
+        node.delete()  # delete the original CCX gate
+
+        return abc.RewriteResult(has_done_something=True)
+
     def rewrite_sxdg(self, node: uop.SXdag) -> abc.RewriteResult:
         return self._rewrite_1q_gates(
             cirq.XPowGate(exponent=-0.5).on(self.cached_qubits[0]), node
@@ -394,9 +441,12 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
         new_gate_stmts = self._generate_1q_gate_stmts(cirq_gate, node.qarg)
         return self._rewrite_gate_stmts(new_gate_stmts, node)
 
-    def _generate_2q_ctrl_gate_stmts(
+    def _generate_multi_ctrl_gate_stmts(
         self, cirq_gate: cirq.Operation, qubits_ssa: List[ir.SSAValue]
     ) -> list[ir.Statement]:
+        qubit_to_ssa_map = {
+            q: ssa for q, ssa in zip(self.cached_qubits[: len(qubits_ssa)], qubits_ssa)
+        }
         target_gates = self.gateset.decompose_to_target_gateset(cirq_gate, 0)
         new_stmts = []
         for new_gate in target_gates:
@@ -412,7 +462,7 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
                 new_stmts.append(phi2_stmt)
                 new_stmts.append(
                     uop.UGate(
-                        qarg=qubits_ssa[new_gate.qubits[0].x],
+                        qarg=qubit_to_ssa_map[new_gate.qubits[0]],
                         theta=phi0_stmt.result,
                         phi=phi1_stmt.result,
                         lam=phi2_stmt.result,
@@ -420,18 +470,31 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
                 )
             else:
                 # 2q
-                new_stmts.append(uop.CZ(ctrl=qubits_ssa[0], qarg=qubits_ssa[1]))
+                new_stmts.append(
+                    uop.CZ(
+                        ctrl=qubit_to_ssa_map[new_gate.qubits[0]],
+                        qarg=qubit_to_ssa_map[new_gate.qubits[1]],
+                    )
+                )
 
         return new_stmts
 
     def _rewrite_2q_ctrl_gates(
         self, cirq_gate: cirq.Operation, node: uop.TwoQubitCtrlGate
     ) -> abc.RewriteResult:
-        new_gate_stmts = self._generate_2q_ctrl_gate_stmts(
+        new_gate_stmts = self._generate_multi_ctrl_gate_stmts(
             cirq_gate, [node.ctrl, node.qarg]
         )
         return self._rewrite_gate_stmts(new_gate_stmts, node)
 
+    def _rewrite_3q_ctrl_gates(
+        self, cirq_gate: cirq.Operation, node: uop.CCX
+    ) -> abc.RewriteResult:
+        new_gate_stmts = self._generate_multi_ctrl_gate_stmts(
+            cirq_gate, [node.ctrl1, node.ctrl2, node.qarg]
+        )
+        return self._rewrite_gate_stmts(new_gate_stmts, node)
+
     def _rewrite_gate_stmts(
         self, new_gate_stmts: list[ir.Statement], node: ir.Statement
     ):