Upgrade to kirin v0.21 (#572)

Co-authored-by: Phillip Weinberg <[email protected]>
Co-authored-by: Dennis Liew <[email protected]>
Co-authored-by: kaihsin <[email protected]>
Co-authored-by: kaihsin <[email protected]>

Author: 4 people

pyproject.toml

@@ -13,7 +13,7 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
     "scipy>=1.13.1",
-    "kirin-toolchain~=0.17.30",
+    "kirin-toolchain~=0.21.0",
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",

src/bloqade/analysis/address/analysis.py

@@ -15,7 +15,7 @@ class AddressAnalysis(Forward[Address]):
     This analysis pass can be used to track the global addresses of qubits and wires.
     """
 
-    keys = ["qubit.address"]
+    keys = ("qubit.address",)
     _const_prop: const.Propagate
     lattice = Address
     next_address: int = field(init=False)
@@ -45,7 +45,7 @@ def try_eval_const_prop(
     ) -> interp.StatementResult[Address]:
         _frame = self._const_prop.initialize_frame(frame.code)
         _frame.set_values(stmt.args, tuple(x.result for x in args))
-        result = self._const_prop.eval_stmt(_frame, stmt)
+        result = self._const_prop.frame_eval(_frame, stmt)
 
         match result:
             case interp.ReturnValue(constant_ret):
@@ -96,7 +96,8 @@ def run_lattice(
         self,
         callee: Address,
         inputs: tuple[Address, ...],
-        kwargs: tuple[str, ...],
+        keys: tuple[str, ...],
+        kwargs: tuple[Address, ...],
     ) -> Address:
         """Run a callable lattice element with the given inputs and keyword arguments.
 
@@ -111,15 +112,16 @@ def run_lattice(
         """
 
         match callee:
-            case PartialLambda(code=code, argnames=argnames):
-                _, ret = self.run_callable(
-                    code, (callee,) + self.permute_values(argnames, inputs, kwargs)
+            case PartialLambda(code=code):
+                _, ret = self.call(
+                    code, callee, *inputs, **{k: v for k, v in zip(keys, kwargs)}
                 )
-                return ret
             case ConstResult(const.Value(ir.Method() as method)):
-                _, ret = self.run_method(
-                    method,
-                    self.permute_values(method.arg_names, inputs, kwargs),
+                _, ret = self.call(
+                    method.code,
+                    self.method_self(method),
+                    *inputs,
+                    **{k: v for k, v in zip(keys, kwargs)},
                 )
                 return ret
             case _:
@@ -137,14 +139,12 @@ def get_const_value(self, addr: Address, typ: Type[T]) -> T | None:
 
         return value
 
-    def eval_stmt_fallback(self, frame: ForwardFrame[Address], stmt: ir.Statement):
-        args = frame.get_values(stmt.args)
+    def eval_fallback(self, frame: ForwardFrame[Address], node: ir.Statement):
+        args = frame.get_values(node.args)
         if types.is_tuple_of(args, ConstResult):
-            return self.try_eval_const_prop(frame, stmt, args)
+            return self.try_eval_const_prop(frame, node, args)
 
-        return tuple(Address.from_type(result.type) for result in stmt.results)
+        return tuple(Address.from_type(result.type) for result in node.results)
 
-    def run_method(self, method: ir.Method, args: tuple[Address, ...]):
-        # NOTE: we do not support dynamic calls here, thus no need to propagate method object
-        self_mt = ConstResult(const.Value(method))
-        return self.run_callable(method.code, (self_mt,) + args)
+    def method_self(self, method: ir.Method) -> Address:
+        return ConstResult(const.Value(method))

src/bloqade/analysis/address/impls.py

@@ -97,7 +97,7 @@ def map_(
 
         results = []
         for ele in values:
-            ret = interp_.run_lattice(fn, (ele,), ())
+            ret = interp_.run_lattice(fn, (ele,), (), ())
             results.append(ret)
 
         if isinstance(stmt, ilist.Map):
@@ -180,13 +180,10 @@ def invoke(
         frame: ForwardFrame[Address],
         stmt: func.Invoke,
     ):
-
-        args = interp_.permute_values(
-            stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
-        )
-        _, ret = interp_.run_method(
-            stmt.callee,
-            args,
+        _, ret = interp_.call(
+            stmt.callee.code,
+            interp_.method_self(stmt.callee),
+            *frame.get_values(stmt.inputs),
         )
 
         return (ret,)
@@ -219,7 +216,8 @@ def call(
         result = interp_.run_lattice(
             frame.get(stmt.callee),
             frame.get_values(stmt.inputs),
-            stmt.kwargs,
+            stmt.keys,
+            frame.get_values(stmt.kwargs),
         )
         return (result,)
 
@@ -319,26 +317,28 @@ def ifelse(
         ):
             body = stmt.then_body if const_cond.data else stmt.else_body
             with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
-                ret = interp_.run_ssacfg_region(body_frame, body, (address_cond,))
+                ret = interp_.frame_call_region(body_frame, stmt, body, address_cond)
                 # interp_.set_values(frame, body_frame.entries.keys(), body_frame.entries.values())
                 return ret
         else:
             # run both branches
             with interp_.new_frame(stmt, has_parent_access=True) as then_frame:
-                then_results = interp_.run_ssacfg_region(
-                    then_frame, stmt.then_body, (address_cond,)
-                )
-                interp_.set_values(
-                    frame, then_frame.entries.keys(), then_frame.entries.values()
+                then_results = interp_.frame_call_region(
+                    then_frame,
+                    stmt,
+                    stmt.then_body,
+                    address_cond,
                 )
+                frame.set_values(then_frame.entries.keys(), then_frame.entries.values())
 
             with interp_.new_frame(stmt, has_parent_access=True) as else_frame:
-                else_results = interp_.run_ssacfg_region(
-                    else_frame, stmt.else_body, (address_cond,)
-                )
-                interp_.set_values(
-                    frame, else_frame.entries.keys(), else_frame.entries.values()
+                else_results = interp_.frame_call_region(
+                    else_frame,
+                    stmt,
+                    stmt.else_body,
+                    address_cond,
                 )
+                frame.set_values(else_frame.entries.keys(), else_frame.entries.values())
             # TODO: pick the non-return value
             if isinstance(then_results, interp.ReturnValue) and isinstance(
                 else_results, interp.ReturnValue
@@ -364,12 +364,12 @@ def for_loop(
         iter_type, iterable = interp_.unpack_iterable(frame.get(stmt.iterable))
 
         if iter_type is None:
-            return interp_.eval_stmt_fallback(frame, stmt)
+            return interp_.eval_fallback(frame, stmt)
 
         for value in iterable:
             with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
-                loop_vars = interp_.run_ssacfg_region(
-                    body_frame, stmt.body, (value,) + loop_vars
+                loop_vars = interp_.frame_call_region(
+                    body_frame, stmt, stmt.body, value, *loop_vars
                 )
 
             if loop_vars is None:

src/bloqade/analysis/fidelity/analysis.py

@@ -4,7 +4,6 @@
 from kirin import ir
 from kirin.lattice import EmptyLattice
 from kirin.analysis import Forward
-from kirin.interp.value import Successor
 from kirin.analysis.forward import ForwardFrame
 
 from ..address import Address, AddressAnalysis
@@ -48,15 +47,11 @@ def main():
     The fidelity of the gate set described by the analysed program. It reduces whenever a noise channel is encountered.
     """
 
-    _current_gate_fidelity: float = field(init=False)
-
     atom_survival_probability: list[float] = field(init=False)
     """
     The probabilities that each of the atoms in the register survive the duration of the analysed program. The order of the list follows the order they are in the register.
     """
 
-    _current_atom_survival_probability: list[float] = field(init=False)
-
     addr_frame: ForwardFrame[Address] = field(init=False)
 
     def initialize(self):
@@ -67,28 +62,21 @@ def initialize(self):
         ]
         return self
 
-    def posthook_succ(self, frame: ForwardFrame, succ: Successor):
-        self.gate_fidelity *= self._current_gate_fidelity
-        for i, _current_survival in enumerate(self._current_atom_survival_probability):
-            self.atom_survival_probability[i] *= _current_survival
-
-    def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
+    def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
         # NOTE: default is to conserve fidelity, so do nothing here
         return
 
-    def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+    def run(self, method: ir.Method, *args, **kwargs) -> tuple[ForwardFrame, Any]:
+        self._run_address_analysis(method)
+        return super().run(method, *args, **kwargs)
 
-    def run_analysis(
-        self, method: ir.Method, args: tuple | None = None, *, no_raise: bool = True
-    ) -> tuple[ForwardFrame, Any]:
-        self._run_address_analysis(method, no_raise=no_raise)
-        return super().run_analysis(method, args, no_raise=no_raise)
-
-    def _run_address_analysis(self, method: ir.Method, no_raise: bool):
+    def _run_address_analysis(self, method: ir.Method):
         addr_analysis = AddressAnalysis(self.dialects)
-        addr_frame, _ = addr_analysis.run_analysis(method=method, no_raise=no_raise)
+        addr_frame, _ = addr_analysis.run(method=method)
         self.addr_frame = addr_frame
 
         # NOTE: make sure we have as many probabilities as we have addresses
         self.atom_survival_probability = [1.0] * addr_analysis.qubit_count
+
+    def method_self(self, method: ir.Method) -> EmptyLattice:
+        return self.lattice.bottom()

src/bloqade/analysis/measure_id/analysis.py

@@ -22,20 +22,16 @@ class MeasurementIDAnalysis(ForwardExtra[MeasureIDFrame, MeasureId]):
     measure_count = 0
 
     def initialize_frame(
-        self, code: ir.Statement, *, has_parent_access: bool = False
+        self, node: ir.Statement, *, has_parent_access: bool = False
     ) -> MeasureIDFrame:
-        return MeasureIDFrame(code, has_parent_access=has_parent_access)
+        return MeasureIDFrame(node, has_parent_access=has_parent_access)
 
     # Still default to bottom,
     # but let constants return the softer "NoMeasureId" type from impl
-    def eval_stmt_fallback(
-        self, frame: ForwardFrame[MeasureId], stmt: ir.Statement
+    def eval_fallback(
+        self, frame: ForwardFrame[MeasureId], node: ir.Statement
     ) -> tuple[MeasureId, ...]:
-        return tuple(NotMeasureId() for _ in stmt.results)
-
-    def run_method(self, method: ir.Method, args: tuple[MeasureId, ...]):
-        # NOTE: we do not support dynamic calls here, thus no need to propagate method object
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+        return tuple(NotMeasureId() for _ in node.results)
 
     # Xiu-zhe (Roger) Luo came up with this in the address analysis,
     # reused here for convenience (now modified to be a bit more graceful)
@@ -45,11 +41,14 @@ def run_method(self, method: ir.Method, args: tuple[MeasureId, ...]):
     T = TypeVar("T")
 
     def get_const_value(
-        self, input_type: type[T], value: ir.SSAValue
+        self, input_type: type[T] | tuple[type[T], ...], value: ir.SSAValue
     ) -> type[T] | None:
         if isinstance(hint := value.hints.get("const"), const.Value):
             data = hint.data
             if isinstance(data, input_type):
                 return hint.data
 
         return None
+
+    def method_self(self, method: ir.Method) -> MeasureId:
+        return self.lattice.bottom()

src/bloqade/analysis/measure_id/impls.py

@@ -152,11 +152,10 @@ def return_(self, _: MeasurementIDAnalysis, frame: interp.Frame, stmt: func.Retu
     def invoke(
         self, interp_: MeasurementIDAnalysis, frame: interp.Frame, stmt: func.Invoke
     ):
-        _, ret = interp_.run_method(
-            stmt.callee,
-            interp_.permute_values(
-                stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
-            ),
+        _, ret = interp_.call(
+            stmt.callee.code,
+            interp_.method_self(stmt.callee),
+            *frame.get_values(stmt.inputs),
         )
         return (ret,)