merging main

Author: weinbe58

pyproject.toml

@@ -17,6 +17,8 @@ dependencies = [
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",
+    "pyqrack>=1.38.2 ; sys_platform == 'darwin'",
+    "pyqrack-cpu>=1.38.2 ; sys_platform != 'darwin'",
 ]
 
 [project.optional-dependencies]
@@ -36,6 +38,12 @@ cirq = [
     "cirq-core>=1.4.1",
     "cirq-core[contrib]>=1.4.1",
 ]
+pyqrack-opencl = [
+    "pyqrack>=1.38.2 ; sys_platform != 'darwin'",
+]
+pyqrack-cuda = [
+    "pyqrack-cuda>=1.38.2",
+]
 
 [build-system]
 requires = ["hatchling"]

src/bloqade/noise/native/stmts.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
 from kirin.dialects import ilist
@@ -7,25 +9,44 @@
 from ._dialect import dialect
 
 
-@statement(dialect=dialect)
-class PauliChannel(ir.Statement):
-
+@statement
+class NativeNoiseStmt(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        raise NotImplementedError(f"Override the method in {type(self).__name__}")
+
+    def check(self):
+        for probs in self.probabilities:
+            self.check_probability(sum(probs))
+            for p in probs:
+                self.check_probability(p)
+
+    def check_probability(self, p: float):
+        if not 0 <= p <= 1:
+            raise ValueError(
+                f"Invalid noise probability encountered in {type(self).__name__}: {p}"
+            )
+
+
+@statement(dialect=dialect)
+class PauliChannel(NativeNoiseStmt):
     px: float = info.attribute(types.Float)
     py: float = info.attribute(types.Float)
     pz: float = info.attribute(types.Float)
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return ((self.px, self.py, self.pz),)
+
 
 NumQubits = types.TypeVar("NumQubits")
 
 
 @statement(dialect=dialect)
-class CZPauliChannel(ir.Statement):
-
-    traits = frozenset({lowering.FromPythonCall()})
-
+class CZPauliChannel(NativeNoiseStmt):
     paired: bool = info.attribute(types.Bool)
     px_ctrl: float = info.attribute(types.Float)
     py_ctrl: float = info.attribute(types.Float)
@@ -36,11 +57,19 @@ class CZPauliChannel(ir.Statement):
     ctrls: ir.SSAValue = info.argument(ilist.IListType[QubitType, NumQubits])
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType, NumQubits])
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return (
+            (self.px_ctrl, self.py_ctrl, self.pz_ctrl),
+            (self.px_qarg, self.py_qarg, self.pz_qarg),
+        )
 
-@statement(dialect=dialect)
-class AtomLossChannel(ir.Statement):
-
-    traits = frozenset({lowering.FromPythonCall()})
 
+@statement(dialect=dialect)
+class AtomLossChannel(NativeNoiseStmt):
     prob: float = info.attribute(types.Float)
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return ((self.prob,),)

src/bloqade/qasm2/dialects/expr/_emit.py

@@ -1,7 +1,6 @@
 from typing import Literal
 
 from kirin import interp
-from kirin.emit.exceptions import EmitError
 
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.types import QubitType
@@ -19,17 +18,18 @@ def emit_func(
         self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.GateFunction
     ):
 
-        args, cparams, qparams = [], [], []
-        for arg in stmt.body.blocks[0].args[1:]:
-            name = frame.get_typed(arg, ast.Name)
-            args.append(name)
-            if not isinstance(name, ast.Name):
-                raise EmitError("expected ast.Name")
+        args: list[ast.Node] = []
+        cparams, qparams = [], []
+        for arg in stmt.body.blocks[0].args:
+            assert arg.name is not None
+
+            args.append(ast.Name(id=arg.name))
             if arg.type.is_subseteq(QubitType):
-                qparams.append(name.id)
+                qparams.append(arg.name)
             else:
-                cparams.append(name.id)
-        emit.run_ssacfg_region(frame, stmt.body, args)
+                cparams.append(arg.name)
+
+        emit.run_ssacfg_region(frame, stmt.body, tuple(args))
         emit.output = ast.Gate(
             name=stmt.sym_name,
             cparams=cparams,

src/bloqade/qasm2/emit/gate.py

@@ -86,21 +86,3 @@ def emit_err(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
     @interp.impl(func.ConstantNone)
     def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         return ()
-
-    @interp.impl(func.Function)
-    def emit_func(
-        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Function
-    ):
-        args_ssa = stmt.args
-        print(stmt.args)
-        emit.run_ssacfg_region(frame, stmt.body, frame.get_values(args_ssa))
-
-        cparams, qparams = [], []
-        for arg in args_ssa:
-            if arg.type.is_subseteq(QubitType):
-                qparams.append(frame.get(arg))
-            else:
-                cparams.append(frame.get(arg))
-
-        emit.output = ast.Gate(stmt.sym_name, cparams, qparams, frame.body)
-        return ()

src/bloqade/qasm2/emit/target.py

@@ -101,7 +101,7 @@ def emit(self, entry: ir.Method) -> ast.MainProgram:
 
         Py2QASM(entry.dialects)(entry)
         target_main = EmitQASM2Main(self.main_target)
-        target_main.run(entry, tuple(ast.Name(name) for name in entry.arg_names[1:]))
+        target_main.run(entry, ())
 
         main_program = target_main.output
         assert main_program is not None, f"failed to emit {entry.sym_name}"

src/bloqade/squin/qubit.py

@@ -117,7 +117,7 @@ def measure(qubit: ilist.IList[Qubit, Any] | list[Qubit]) -> list[bool]: ...
 
 @wraps(MeasureAny)
 def measure(qubit: Any) -> Any:
-    """Measure a qubit or qubits in the list."
+    """Measure a qubit or qubits in the list.
 
     Args:
         qubits: The list of qubits to measure.
@@ -128,6 +128,22 @@ def measure(qubit: Any) -> Any:
     ...
 
 
+@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:
+        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
+    """
+    ...
+
+
 @wraps(MeasureAndReset)
 def measure_and_reset(qubits: ilist.IList[Qubit, Any]) -> int:
     """Measure the qubits in the list and reset them."

test/pyqrack/runtime/noise/native/test_pauli.py

@@ -1,5 +1,6 @@
 from unittest.mock import Mock, call
 
+import pytest
 from kirin import ir
 
 from bloqade import qasm2
@@ -41,6 +42,23 @@ def test_atom_loss():
     sim_reg.assert_has_calls([call.y(0)])
 
 
+@pytest.mark.xfail
+def test_pauli_probs_check():
+    @simulation
+    def test_atom_loss():
+        q = qasm2.qreg(2)
+        native.pauli_channel(
+            [q[0]],
+            px=0.1,
+            py=0.4,
+            pz=1.3,
+        )
+        return q
+
+    with pytest.raises(ir.ValidationError):
+        test_atom_loss.verify()
+
+
 def test_cz_pauli_channel_false():
     @simulation
     def test_atom_loss():
@@ -122,9 +140,3 @@ def test_atom_loss():
     sim_reg = run_mock(test_atom_loss, rng_state)
 
     sim_reg.assert_has_calls([call.y(0), call.x(1), call.mcz([0], 1)])
-
-
-if __name__ == "__main__":
-    test_pauli_channel()
-    test_cz_pauli_channel_false()
-    test_cz_pauli_channel_true()

test/qasm2/emit/test_qasm2.py

@@ -1,9 +1,6 @@
-import pytest
-
 from bloqade import qasm2
 
 
-@pytest.mark.skip(reason="broken gate emit!")
 def test_qasm2_custom_gate():
     @qasm2.gate
     def custom_gate(a: qasm2.Qubit, b: qasm2.Qubit):