Merge branch 'main' into john/add-broadcast-to-wire

Author: weinbe58

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v6
         with:
           # Install a specific version of uv.
           version: "0.5.1"

.github/workflows/release.yml

@@ -10,7 +10,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Install uv
-      uses: astral-sh/setup-uv@v5
+      uses: astral-sh/setup-uv@v6
       with:
         # Install a specific version of uv.
         version: "0.5.5"

src/bloqade/pyqrack/qasm2/uop.py

@@ -2,6 +2,7 @@
 
 from kirin import interp
 
+from pyqrack.pauli import Pauli
 from bloqade.pyqrack.reg import PyQrackQubit
 from bloqade.qasm2.dialects import uop
 
@@ -26,7 +27,14 @@ class PyQrackMethods(interp.MethodTable):
         "tdg": "adjt",
     }
 
-    AXIS_MAP = {"rx": 1, "ry": 2, "rz": 3, "crx": 1, "cry": 2, "crz": 3}
+    AXIS_MAP = {
+        "rx": Pauli.PauliX,
+        "ry": Pauli.PauliY,
+        "rz": Pauli.PauliZ,
+        "crx": Pauli.PauliX,
+        "cry": Pauli.PauliY,
+        "crz": Pauli.PauliZ,
+    }
 
     @interp.impl(uop.Barrier)
     def barrier(

src/bloqade/qasm2/passes/noise.py

@@ -1,7 +1,7 @@
 from dataclasses import field, dataclass
 
 from kirin import ir
-from kirin.passes import Pass
+from kirin.passes import Pass, HintConst
 from kirin.rewrite import (
     Walk,
     Chain,
@@ -10,11 +10,10 @@
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
-from kirin.rewrite.abc import RewriteResult
 
 from bloqade.noise import native
 from bloqade.analysis import address
-from bloqade.qasm2.rewrite.heuristic_noise import NoiseRewriteRule
+from bloqade.qasm2.rewrite.heuristic_noise import InsertGetQubit, NoiseRewriteRule
 
 
 @dataclass
@@ -38,16 +37,17 @@ 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(InsertGetQubit()).rewrite(mt.code)
+        HintConst(self.dialects).unsafe_run(mt)
+        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
         result = (
             Walk(
                 NoiseRewriteRule(
                     address_analysis=frame.entries,
                     noise_model=self.noise_model,
                     gate_noise_params=self.gate_noise_params,
-                )
+                ),
+                reverse=True,
             )
             .rewrite(mt.code)
             .join(result)

src/bloqade/qasm2/rewrite/heuristic_noise.py

@@ -10,6 +10,34 @@
 from bloqade.qasm2.dialects import uop, core, glob, parallel
 
 
+class InsertGetQubit(rewrite_abc.RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+        if (
+            not isinstance(node, core.QRegNew)
+            or not isinstance(n_qubits_stmt := node.n_qubits.owner, py.Constant)
+            or not isinstance(n_qubits := n_qubits_stmt.value.unwrap(), int)
+            or (block := node.parent_block) is None
+        ):
+            return rewrite_abc.RewriteResult()
+
+        n_qubits_stmt.detach()
+        node.detach()
+        if block.first_stmt is None:
+            block.stmts.append(node)
+        else:
+            node.insert_before(block.first_stmt)
+            n_qubits_stmt.insert_before(block.first_stmt)
+
+        for idx_val in range(n_qubits):
+            idx = py.constant.Constant(value=idx_val)
+            qubit = core.QRegGet(node.result, idx=idx.result)
+            qubit.insert_after(node)
+            idx.insert_after(node)
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+
 @dataclass
 class NoiseRewriteRule(rewrite_abc.RewriteRule):
     """
@@ -24,12 +52,18 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
     noise_model: native.MoveNoiseModelABC = field(
         default_factory=native.TwoRowZoneModel
     )
-    qubit_ssa_value: Dict[int, ir.SSAValue] = field(default_factory=dict, init=False)
+
+    def __post_init__(self):
+        self.qubit_ssa_value: Dict[int, ir.SSAValue] = {}
+        for ssa_value, addr in self.address_analysis.items():
+            if (
+                isinstance(addr, address.AddressQubit)
+                and ssa_value not in self.qubit_ssa_value
+            ):
+                self.qubit_ssa_value[addr.data] = ssa_value
 
     def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
-        if isinstance(node, core.QRegNew):
-            return self.rewrite_qreg_new(node)
-        elif isinstance(node, uop.SingleQubitGate):
+        if isinstance(node, uop.SingleQubitGate):
             return self.rewrite_single_qubit_gate(node)
         elif isinstance(node, uop.CZ):
             return self.rewrite_cz_gate(node)
@@ -42,24 +76,6 @@ def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
         else:
             return rewrite_abc.RewriteResult()
 
-    def rewrite_qreg_new(self, node: core.QRegNew):
-
-        addr = self.address_analysis[node.result]
-        if not isinstance(addr, address.AddressReg):
-            return rewrite_abc.RewriteResult()
-
-        has_done_something = False
-        for idx_val, qid in enumerate(addr.data):
-            if qid not in self.qubit_ssa_value:
-                has_done_something = True
-                idx = py.constant.Constant(value=idx_val)
-                qubit = core.QRegGet(node.result, idx=idx.result)
-                self.qubit_ssa_value[qid] = qubit.result
-                qubit.insert_after(node)
-                idx.insert_after(node)
-
-        return rewrite_abc.RewriteResult(has_done_something=has_done_something)
-
     def insert_single_qubit_noise(
         self,
         node: ir.Statement,

src/bloqade/qasm2/rewrite/uop_to_parallel.py

@@ -154,7 +154,7 @@ def __call__(self, node: ir.Statement) -> RewriteResult:
             self.group_has_merged[group_number] = result.has_done_something
             return result
 
-        if self.group_has_merged[group_number]:
+        if self.group_has_merged.setdefault(group_number, False):
             node.delete()
 
         return RewriteResult(has_done_something=self.group_has_merged[group_number])

src/bloqade/squin/groups.py

@@ -1,15 +1,18 @@
 from kirin import ir, passes
 from kirin.prelude import structural_no_opt
 from kirin.dialects import ilist
+from kirin.rewrite.walk import Walk
 
 from . import op, wire, qubit
+from .rewrite.measure_desugar import MeasureDesugarRule
 
 
 @ir.dialect_group(structural_no_opt.union([op, qubit]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
     ilist_desugar_pass = ilist.IListDesugar(self)
+    measure_desugar_pass = Walk(MeasureDesugarRule())
 
     def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
         method.verify()
@@ -18,6 +21,7 @@ def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
 
         if typeinfer:
             typeinfer_pass(method)
+            measure_desugar_pass.rewrite(method.code)
         ilist_desugar_pass(method)
         if typeinfer:
             typeinfer_pass(method)  # fix types after desugaring

src/bloqade/squin/qubit.py

@@ -7,7 +7,7 @@
 - `kirin.dialects.ilist`: provides the `ilist.IListType` type for lists of qubits.
 """
 
-from typing import Any
+from typing import Any, overload
 
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
@@ -41,8 +41,26 @@ class New(ir.Statement):
     result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
 
 
+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 Measure(ir.Statement):
+class MeasureQubit(ir.Statement):
+    name = "measure.qubit"
+
+    traits = frozenset({lowering.FromPythonCall()})
+    qubit: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
+
+
+@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[types.Bool])
@@ -89,31 +107,39 @@ def apply(operator: Op, qubits: ilist.IList[Qubit, Any] | list[Qubit]) -> None:
     ...
 
 
-@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.
+@overload
+def measure(input: Qubit) -> bool: ...
+@overload
+def measure(input: ilist.IList[Qubit, Any] | list[Qubit]) -> list[bool]: ...
+
+
+@wraps(MeasureAny)
+def measure(input: Any) -> Any:
+    """Measure a qubit or qubits in the list.
 
     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.
+        input: A qubit or a list of qubits to measure.
 
     Returns:
-        None
+        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(Measure)
-def measure(qubits: ilist.IList[Qubit, Any]) -> int:
-    """Measure the qubits in the list."
+@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.
 
     Args:
-        qubits: The list of qubits to measure.
+        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:
-        int: The result of the measurement.
+        None
     """
     ...
 

src/bloqade/squin/rewrite/__init__.py

@@ -0,0 +1,33 @@
+from kirin import ir, types
+from kirin.dialects import ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin.qubit import 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()