upgrade to kirin 0.17 (#178)

Author: Roger-luo

pyproject.toml

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

src/bloqade/noise/native/rewrite.py

@@ -1,27 +1,27 @@
 from dataclasses import dataclass
 
 from kirin import ir
-from kirin.rewrite import abc, dce, walk, result, fixpoint
+from kirin.rewrite import abc, dce, walk, fixpoint
 from kirin.passes.abc import Pass
 
 from .stmts import PauliChannel, CZPauliChannel, AtomLossChannel
 from ._dialect import dialect
 
 
 class RemoveNoiseRewrite(abc.RewriteRule):
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, (AtomLossChannel, PauliChannel, CZPauliChannel)):
             node.delete()
-            return result.RewriteResult(has_done_something=True)
+            return abc.RewriteResult(has_done_something=True)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
 
 @dataclass
 class RemoveNoisePass(Pass):
     name = "remove-noise"
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         delete_walk = walk.Walk(RemoveNoiseRewrite())
         dce_walk = fixpoint.Fixpoint(walk.Walk(dce.DeadCodeElimination()))
 

src/bloqade/pyqrack/target.py

@@ -80,7 +80,7 @@ def run(
         """
         fold = Fold(mt.dialects)
         fold(mt)
-        return self._get_interp(mt).run(mt, args, kwargs).expect()
+        return self._get_interp(mt).run(mt, args, kwargs)
 
     def multi_run(
         self,
@@ -107,6 +107,6 @@ def multi_run(
         interpreter = self._get_interp(mt)
         batched_results = []
         for _ in range(_shots):
-            batched_results.append(interpreter.run(mt, args, kwargs).expect())
+            batched_results.append(interpreter.run(mt, args, kwargs))
 
         return batched_results

src/bloqade/qasm2/_wrappers.py

@@ -61,33 +61,25 @@ def reset(qarg: Qubit) -> None:
 
 
 @overload
-def measure(qarg: Qubit, cbit: Bit) -> None:
-    """
-    Measure the qubit `qarg` and store the result in the classical bit `cbit`.
-
-    Args:
-        qarg: The qubit to measure.
-        cbit: The classical bit to store the result in.
-    """
-    ...
+def measure(qreg: QReg, creg: CReg) -> None: ...
 
 
 @overload
-def measure(qarg: QReg, carg: CReg) -> None:
+def measure(qarg: Qubit, cbit: Bit) -> None: ...
+
+
+@wraps(core.Measure)
+def measure(qarg, cbit) -> None:
     """
-    Measure each qubit in the quantum register `qarg` and store the result in the classical register `carg`.
+    Measure the qubit `qarg` and store the result in the classical bit `cbit`.
 
     Args:
-        qarg: The quantum register to measure.
-        carg: The classical bit to store the result in.
+        qarg: The qubit to measure.
+        cbit: The classical bit to store the result in.
     """
     ...
 
 
-@wraps(core.Measure)
-def measure(qarg, carg) -> None: ...
-
-
 @wraps(uop.CX)
 def cx(ctrl: Qubit, qarg: Qubit) -> None:
     """

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

@@ -19,16 +19,17 @@ def emit_func(
         self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.GateFunction
     ):
 
-        cparams, qparams = [], []
+        args, cparams, qparams = [], [], []
         for arg in stmt.body.blocks[0].args[1:]:
-            name = frame.get(arg)
+            name = frame.get_typed(arg, ast.Name)
+            args.append(name)
             if not isinstance(name, ast.Name):
                 raise EmitError("expected ast.Name")
             if arg.type.is_subseteq(QubitType):
                 qparams.append(name.id)
             else:
                 cparams.append(name.id)
-        emit.run_ssacfg_region(frame, stmt.body)
+        emit.run_ssacfg_region(frame, stmt.body, args)
         emit.output = ast.Gate(
             name=stmt.sym_name,
             cparams=cparams,

src/bloqade/qasm2/emit/base.py

@@ -36,8 +36,10 @@ def initialize(self) -> Self:
         )
         return self
 
-    def new_frame(self, code: ir.Statement) -> EmitQASM2Frame:
-        return EmitQASM2Frame.from_func_like(code)
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> EmitQASM2Frame[StmtType]:
+        return EmitQASM2Frame(code, has_parent_access=has_parent_access)
 
     def run_method(
         self, method: ir.Method, args: tuple[ast.Node | None, ...]

src/bloqade/qasm2/emit/gate.py

@@ -91,12 +91,16 @@ def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
     def emit_func(
         self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Function
     ):
-        emit.run_ssacfg_region(frame, stmt.body)
+        args_ssa = stmt.args
+        print(stmt.args)
+        emit.run_ssacfg_region(frame, stmt.body, frame.get_values(args_ssa))
+
         cparams, qparams = [], []
-        for arg in stmt.args:
+        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/main.py

@@ -24,7 +24,7 @@ def emit_func(
     ):
         from bloqade.qasm2.dialects import glob, noise, parallel
 
-        emit.run_ssacfg_region(frame, stmt.body)
+        emit.run_ssacfg_region(frame, stmt.body, ())
         if emit.dialects.data.intersection(
             (parallel.dialect, glob.dialect, noise.dialect)
         ):
@@ -51,12 +51,14 @@ def emit_conditional_branch(
         self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: cf.ConditionalBranch
     ):
         cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
-        body_frame = emit.new_frame(stmt)
-        body_frame.entries.update(frame.entries)
-        body_frame.set_values(
-            stmt.then_successor.args, frame.get_values(stmt.then_arguments)
-        )
-        emit.emit_block(body_frame, stmt.then_successor)
+
+        with emit.new_frame(stmt) as body_frame:
+            body_frame.entries.update(frame.entries)
+            body_frame.set_values(
+                stmt.then_successor.args, frame.get_values(stmt.then_arguments)
+            )
+            emit.emit_block(body_frame, stmt.then_successor)
+
         frame.body.append(
             ast.IfStmt(
                 cond,
@@ -91,15 +93,17 @@ def emit_if_else(
             )
 
         cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
-        then_frame = emit.new_frame(stmt)
-        then_frame.entries.update(frame.entries)
-        emit.emit_block(then_frame, stmt.then_body.blocks[0])
-        frame.body.append(
-            ast.IfStmt(
-                cond,
-                body=then_frame.body,  # type: ignore
+
+        with emit.new_frame(stmt) as then_frame:
+            then_frame.entries.update(frame.entries)
+            emit.emit_block(then_frame, stmt.then_body.blocks[0])
+            frame.body.append(
+                ast.IfStmt(
+                    cond,
+                    body=then_frame.body,  # type: ignore
+                )
             )
-        )
+
         term = stmt.then_body.blocks[0].last_stmt
         if isinstance(term, scf.Yield):
             return then_frame.get_values(term.values)

src/bloqade/qasm2/emit/target.py

@@ -101,9 +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:])
-        ).expect()
+        target_main.run(entry, tuple(ast.Name(name) for name in entry.arg_names[1:]))
 
         main_program = target_main.output
         assert main_program is not None, f"failed to emit {entry.sym_name}"
@@ -133,9 +131,7 @@ def emit(self, entry: ir.Method) -> ast.MainProgram:
 
                 Py2QASM(fn.dialects)(fn)
 
-                target_gate.run(
-                    fn, tuple(ast.Name(name) for name in fn.arg_names[1:])
-                ).expect()
+                target_gate.run(fn, tuple(ast.Name(name) for name in fn.arg_names[1:]))
                 assert target_gate.output is not None, f"failed to emit {fn.sym_name}"
                 extra.append(target_gate.output)
 

src/bloqade/qasm2/passes/fold.py

@@ -19,7 +19,7 @@
 from kirin.analysis import const
 from kirin.dialects import scf, ilist
 from kirin.ir.method import Method
-from kirin.rewrite.result import RewriteResult
+from kirin.rewrite.abc import RewriteResult
 
 from bloqade.qasm2.dialects import expr