Merge branch 'main' into roger/upgrade-kirin-0-17

Author: Roger-luo

src/bloqade/analysis/address/impls.py

@@ -192,7 +192,10 @@ def unwrap(
 
         origin_qubit = frame.get(stmt.qubit)
 
-        return (AddressWire(origin_qubit=origin_qubit),)
+        if isinstance(origin_qubit, AddressQubit):
+            return (AddressWire(origin_qubit=origin_qubit),)
+        else:
+            return (Address.top(),)
 
     @interp.impl(squin.wire.Apply)
     def apply(
@@ -201,14 +204,7 @@ def apply(
         frame: ForwardFrame[Address],
         stmt: squin.wire.Apply,
     ):
-
-        origin_qubits = tuple(
-            [frame.get(input_elem).origin_qubit for input_elem in stmt.inputs]
-        )
-        new_address_wires = tuple(
-            [AddressWire(origin_qubit=origin_qubit) for origin_qubit in origin_qubits]
-        )
-        return new_address_wires
+        return frame.get_values(stmt.inputs)
 
 
 @squin.qubit.dialect.register(key="qubit.address")

src/bloqade/analysis/address/lattice.py

@@ -81,5 +81,5 @@ class AddressWire(Address):
 
     def is_subseteq(self, other: Address) -> bool:
         if isinstance(other, AddressWire):
-            return self.origin_qubit == self.origin_qubit
+            return self.origin_qubit == other.origin_qubit
         return False

src/bloqade/noise/native/model.py

@@ -102,10 +102,9 @@ class MoveNoiseModelABC(abc.ABC):
     params: MoveNoiseParams = field(default_factory=MoveNoiseParams)
     """Parameters for calculating move noise."""
 
-    @classmethod
     @abc.abstractmethod
     def parallel_cz_errors(
-        cls, ctrls: List[int], qargs: List[int], rest: List[int]
+        self, ctrls: List[int], qargs: List[int], rest: List[int]
     ) -> Dict[Tuple[float, float, float, float], List[int]]:
         """Takes a set of ctrls and qargs and returns a noise model for all qubits."""
         pass

src/bloqade/pyqrack/__init__.py

@@ -14,5 +14,5 @@
 
 # NOTE: The following import is for registering the method tables
 from .noise import native as native
-from .qasm2 import uop as uop, core as core, parallel as parallel
+from .qasm2 import uop as uop, core as core, glob as glob, parallel as parallel
 from .target import PyQrack as PyQrack

src/bloqade/pyqrack/noise/native.py

@@ -93,8 +93,7 @@ def atom_loss_channel(
 
         for qarg in active_qubits:
             if interp.rng_state.uniform() <= stmt.prob:
-                sim_reg = qarg.ref.sim_reg
-                sim_reg.force_m(qarg.addr, 0)
+                qarg.ref.sim_reg.m(qarg.addr)
                 qarg.drop()
 
         return ()

src/bloqade/pyqrack/qasm2/core.py

@@ -14,7 +14,6 @@
 
 @core.dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
-
     @interp.impl(core.QRegNew)
     def qreg_new(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.QRegNew

src/bloqade/pyqrack/qasm2/glob.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.pyqrack.reg import PyQrackReg
+from bloqade.pyqrack.base import PyQrackInterpreter
+from bloqade.qasm2.dialects import glob
+
+
+@glob.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    @interp.impl(glob.UGate)
+    def ugate(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: glob.UGate):
+        registers: ilist.IList[PyQrackReg, Any] = frame.get(stmt.registers)
+        theta, phi, lam = (
+            frame.get(stmt.theta),
+            frame.get(stmt.phi),
+            frame.get(stmt.lam),
+        )
+
+        for qreg in registers:
+            for qarg in qreg:
+                if qarg.is_active():
+                    interp.memory.sim_reg.u(qarg.addr, theta, phi, lam)
+        return ()

src/bloqade/pyqrack/reg.py

@@ -33,7 +33,7 @@ class CBitRef:
     pos: int
     """The position of this bit in the classical register."""
 
-    def set_value(self, value: bool):
+    def set_value(self, value: Measurement):
         self.ref[self.pos] = value
 
     def get_value(self):
@@ -46,7 +46,7 @@ class QubitState(enum.Enum):
 
 
 @dataclass(frozen=True)
-class PyQrackReg(QReg):
+class PyQrackReg(QReg):  # TODO: clean up implementation with list base class
     """Simulation runtime value of a quantum register."""
 
     size: int
@@ -72,6 +72,8 @@ def drop(self, pos: int):
         self.qubit_state[pos] = QubitState.Lost
 
     def __getitem__(self, pos: int):
+        if not 0 <= pos < self.size:
+            raise IndexError("Qubit index out of bounds of register.")
         return PyQrackQubit(self, pos)
 
 

src/bloqade/qasm2/glob.py

@@ -11,7 +11,7 @@
 
 @wraps(glob.UGate)
 def u(
-    theta: float, phi: float, lam: float, registers: ilist.IList[QReg, Any] | list
+    registers: ilist.IList[QReg, Any] | list, theta: float, phi: float, lam: float
 ) -> None:
     """Apply a U gate to all qubits in the input registers.
 

src/bloqade/qasm2/parse/lowering.py

@@ -1,3 +1,5 @@
+import os
+import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -17,6 +19,118 @@ 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(
+        self,
+        source: str,
+        kernel_name: str,
+        *,
+        returns: 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)
+
+        state = lowering.State(
+            self,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+        with state.frame(
+            [stmt],
+            globals=globals,
+        ) 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()
+
+                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,
+        )
+
+    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(
         self,
         stmt: ast.Node,
@@ -85,6 +199,10 @@ def lower_literal(self, state: lowering.State[ast.Node], value) -> ir.SSAValue:
             stmt = expr.ConstInt(value=value)
         elif isinstance(value, float):
             stmt = expr.ConstFloat(value=value)
+        else:
+            raise lowering.BuildError(
+                f"Expected value of type float or int, got {type(value)}."
+            )
         state.current_frame.push(stmt)
         return stmt.result
 
@@ -99,6 +217,8 @@ def visit_MainProgram(self, state: lowering.State[ast.Node], node: ast.MainProgr
             dialects = ["qasm2.core", "qasm2.uop", "qasm2.expr"]
         elif isinstance(node.header, ast.Kirin):
             dialects = node.header.dialects
+        else:
+            raise lowering.BuildError(f"Unexpected node header {node.header}")
 
         for dialect in dialects:
             if dialect not in allowed:
@@ -295,6 +415,8 @@ def visit_Bit(self, state: lowering.State[ast.Node], node: ast.Bit):
             stmt = core.QRegGet(reg, addr.result)
         elif reg.type.is_subseteq(CRegType):
             stmt = core.CRegGet(reg, addr.result)
+        else:
+            raise lowering.BuildError(f"Unexpected register type {reg.type}")
         return state.current_frame.push(stmt).result
 
     def visit_Call(self, state: lowering.State[ast.Node], node: ast.Call):