upgrade to kirin 0.17

Roger-luo · Roger-luo · commit dafd4e0c9fb4 · 2025-04-21T13:51:04.000-04:00
diff --git a/src/bloqade/noise/native/rewrite.py b/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()))
 
diff --git a/src/bloqade/qasm2/_wrappers.py b/src/bloqade/qasm2/_wrappers.py
@@ -1,3 +1,5 @@
+from typing import overload
+
 from kirin.lowering import wraps
 
 from .types import Bit, CReg, QReg, Qubit
@@ -58,8 +60,16 @@ def reset(qarg: Qubit) -> None:
     ...
 
 
+@overload
+def measure(qreg: QReg, creg: CReg) -> None: ...
+
+
+@overload
+def measure(qarg: Qubit, cbit: Bit) -> None: ...
+
+
 @wraps(core.Measure)
-def measure(qarg: Qubit, cbit: Bit) -> None:
+def measure(qarg, cbit) -> None:
     """
     Measure the qubit `qarg` and store the result in the classical bit `cbit`.
 
diff --git a/src/bloqade/qasm2/dialects/expr/_emit.py b/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,
diff --git a/src/bloqade/qasm2/emit/base.py b/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, ...]
diff --git a/src/bloqade/qasm2/emit/target.py b/src/bloqade/qasm2/emit/target.py
@@ -133,9 +133,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)
 
diff --git a/src/bloqade/qasm2/passes/glob.py b/src/bloqade/qasm2/passes/glob.py
@@ -4,7 +4,7 @@
 """
 
 from kirin import ir
-from kirin.rewrite import cse, dce, walk, result
+from kirin.rewrite import abc, cse, dce, walk
 from kirin.passes.abc import Pass
 from kirin.passes.fold import Fold
 from kirin.rewrite.fixpoint import Fixpoint
@@ -54,7 +54,7 @@ def generate_rule(self, mt: ir.Method) -> GlobalToUOpRule:
         frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
         return GlobalToUOpRule(frame.entries)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         rewriter = walk.Walk(self.generate_rule(mt))
         result = rewriter.rewrite(mt.code)
 
@@ -106,7 +106,7 @@ def generate_rule(self, mt: ir.Method) -> GlobalToParallelRule:
         frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
         return GlobalToParallelRule(frame.entries)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         rewriter = walk.Walk(self.generate_rule(mt))
         result = rewriter.rewrite(mt.code)
 
diff --git a/src/bloqade/qasm2/passes/noise.py b/src/bloqade/qasm2/passes/noise.py
@@ -10,7 +10,7 @@
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
-from kirin.rewrite.result import RewriteResult
+from kirin.rewrite.abc import RewriteResult
 
 from bloqade.noise import native
 from bloqade.analysis import address
diff --git a/src/bloqade/qasm2/passes/parallel.py b/src/bloqade/qasm2/passes/parallel.py
@@ -16,7 +16,7 @@
     ConstantFold,
     DeadCodeElimination,
     CommonSubexpressionElimination,
-    result,
+    abc,
 )
 from kirin.analysis import const
 
@@ -84,7 +84,7 @@ def generate_rule(self, mt: ir.Method) -> ParallelToUOpRule:
 
         return ParallelToUOpRule(id_map=id_map, address_analysis=frame.entries)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         result = Walk(self.generate_rule(mt)).rewrite(mt.code)
         rule = Chain(
             ConstantFold(),
@@ -140,7 +140,7 @@ def test():
     def __post_init__(self):
         self.constprop = const.Propagate(self.dialects)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         result = Walk(RaiseRegisterRule()).rewrite(mt.code)
 
         # do not run the parallelization because registers are not at the top
diff --git a/src/bloqade/qasm2/passes/py2qasm.py b/src/bloqade/qasm2/passes/py2qasm.py
@@ -4,8 +4,7 @@
 from kirin.passes import Pass
 from kirin.rewrite import Walk, Fixpoint
 from kirin.dialects import py, math
-from kirin.rewrite.abc import RewriteRule
-from kirin.rewrite.result import RewriteResult
+from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.qasm2.dialects import core, expr
 
diff --git a/src/bloqade/qasm2/passes/qasm2py.py b/src/bloqade/qasm2/passes/qasm2py.py
@@ -6,8 +6,7 @@
 from kirin.passes import Pass
 from kirin.rewrite import Walk, Fixpoint
 from kirin.dialects import py, math
-from kirin.rewrite.abc import RewriteRule
-from kirin.rewrite.result import RewriteResult
+from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.qasm2.dialects import core, expr
 
diff --git a/src/bloqade/qasm2/rewrite/desugar.py b/src/bloqade/qasm2/rewrite/desugar.py
@@ -2,27 +2,27 @@
 
 from kirin import ir
 from kirin.passes import Pass
-from kirin.rewrite import abc, walk, result
+from kirin.rewrite import abc, walk
 from kirin.dialects import py
 
 from bloqade.qasm2.dialects import core
 
 
 class IndexingDesugarRule(abc.RewriteRule):
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, py.indexing.GetItem):
             if node.obj.type.is_subseteq(core.QRegType):
                 node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
-                return result.RewriteResult(has_done_something=True)
+                return abc.RewriteResult(has_done_something=True)
             elif node.obj.type.is_subseteq(core.CRegType):
                 node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
-                return result.RewriteResult(has_done_something=True)
+                return abc.RewriteResult(has_done_something=True)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
 
 @dataclass
 class IndexingDesugarPass(Pass):
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
 
         return walk.Walk(IndexingDesugarRule()).rewrite(mt.code)
diff --git a/src/bloqade/qasm2/rewrite/glob.py b/src/bloqade/qasm2/rewrite/glob.py
@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 
 from kirin import ir
-from kirin.rewrite import abc, result
+from kirin.rewrite import abc
 from kirin.dialects import py, ilist
 
 from bloqade import qasm2
@@ -47,18 +47,18 @@ def get_qubit_ssa(self, node: glob.UGate):
 @dataclass
 class GlobalToParallelRule(abc.RewriteRule, GlobalRewriteBase):
 
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if type(node) in glob.dialect.stmts:
             return getattr(self, f"rewrite_{node.name}")(node)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
     def rewrite_ugate(self, node: glob.UGate):
 
         new_stmts, qubit_ssa = self.get_qubit_ssa(node)
 
         if qubit_ssa is None:
-            return result.RewriteResult()
+            return abc.RewriteResult()
 
         new_stmts.append(qargs := ilist.New(values=qubit_ssa))
         new_stmts.append(
@@ -72,24 +72,24 @@ def rewrite_ugate(self, node: glob.UGate):
 
         node.delete()
 
-        return result.RewriteResult(has_done_something=True)
+        return abc.RewriteResult(has_done_something=True)
 
 
 @dataclass
 class GlobalToUOpRule(abc.RewriteRule, GlobalRewriteBase):
 
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if type(node) in glob.dialect.stmts:
             return getattr(self, f"rewrite_{node.name}")(node)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
     def rewrite_ugate(self, node: glob.UGate):
 
         new_stmts, qubit_ssa = self.get_qubit_ssa(node)
 
         if qubit_ssa is None:
-            return result.RewriteResult()
+            return abc.RewriteResult()
 
         for qarg in qubit_ssa:
             new_stmts.append(
@@ -100,4 +100,4 @@ def rewrite_ugate(self, node: glob.UGate):
             stmt.insert_before(node)
 
         node.delete()
-        return result.RewriteResult(has_done_something=True)
+        return abc.RewriteResult(has_done_something=True)
diff --git a/src/bloqade/qasm2/rewrite/heuristic_noise.py b/src/bloqade/qasm2/rewrite/heuristic_noise.py
@@ -2,7 +2,7 @@
 from dataclasses import field, dataclass
 
 from kirin import ir
-from kirin.rewrite import abc as result_abc, result
+from kirin.rewrite import abc as rewrite_abc
 from kirin.dialects import py, ilist
 
 from bloqade.noise import native
@@ -11,7 +11,7 @@
 
 
 @dataclass
-class NoiseRewriteRule(result_abc.RewriteRule):
+class NoiseRewriteRule(rewrite_abc.RewriteRule):
     """
     NOTE: This pass is not guaranteed to be supported long-term in bloqade. We will be
     moving towards a more general approach to noise modeling in the future.
@@ -26,7 +26,7 @@ class NoiseRewriteRule(result_abc.RewriteRule):
     )
     qubit_ssa_value: Dict[int, ir.SSAValue] = field(default_factory=dict, init=False)
 
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    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):
@@ -40,13 +40,13 @@ def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
         elif isinstance(node, glob.UGate):
             return self.rewrite_global_single_qubit_gate(node)
         else:
-            return result.RewriteResult()
+            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 result.RewriteResult()
+            return rewrite_abc.RewriteResult()
 
         has_done_something = False
         for idx_val, qid in enumerate(addr.data):
@@ -58,7 +58,7 @@ def rewrite_qreg_new(self, node: core.QRegNew):
                 qubit.insert_after(node)
                 idx.insert_after(node)
 
-        return result.RewriteResult(has_done_something=has_done_something)
+        return rewrite_abc.RewriteResult(has_done_something=has_done_something)
 
     def insert_single_qubit_noise(
         self,
@@ -71,7 +71,7 @@ def insert_single_qubit_noise(
         )
         native.AtomLossChannel(qargs, prob=probs[3]).insert_before(node)
 
-        return result.RewriteResult(has_done_something=True)
+        return rewrite_abc.RewriteResult(has_done_something=True)
 
     def rewrite_single_qubit_gate(self, node: uop.SingleQubitGate):
         probs = (
@@ -86,13 +86,13 @@ def rewrite_single_qubit_gate(self, node: uop.SingleQubitGate):
     def rewrite_global_single_qubit_gate(self, node: glob.UGate):
         addrs = self.address_analysis[node.registers]
         if not isinstance(addrs, address.AddressTuple):
-            return result.RewriteResult()
+            return rewrite_abc.RewriteResult()
 
         qargs = []
 
         for addr in addrs.data:
             if not isinstance(addr, address.AddressReg):
-                return result.RewriteResult()
+                return rewrite_abc.RewriteResult()
 
             for qid in addr.data:
                 qargs.append(self.qubit_ssa_value[qid])
@@ -109,10 +109,10 @@ def rewrite_global_single_qubit_gate(self, node: glob.UGate):
     def rewrite_parallel_single_qubit_gate(self, node: parallel.RZ | parallel.UGate):
         addrs = self.address_analysis[node.qargs]
         if not isinstance(addrs, address.AddressTuple):
-            return result.RewriteResult()
+            return rewrite_abc.RewriteResult()
 
         if not all(isinstance(addr, address.AddressQubit) for addr in addrs.data):
-            return result.RewriteResult()
+            return rewrite_abc.RewriteResult()
 
         probs = (
             self.gate_noise_params.local_px,
@@ -213,7 +213,7 @@ def rewrite_cz_gate(self, node: uop.CZ):
             new_node.insert_before(node)
             has_done_something = True
 
-        return result.RewriteResult(has_done_something=has_done_something)
+        return rewrite_abc.RewriteResult(has_done_something=has_done_something)
 
     def rewrite_parallel_cz_gate(self, node: parallel.CZ):
         ctrls = self.address_analysis[node.ctrls]
@@ -244,4 +244,4 @@ def rewrite_parallel_cz_gate(self, node: parallel.CZ):
             new_node.insert_before(node)
             has_done_something = True
 
-        return result.RewriteResult(has_done_something=has_done_something)
+        return rewrite_abc.RewriteResult(has_done_something=has_done_something)
diff --git a/src/bloqade/qasm2/rewrite/native_gates.py b/src/bloqade/qasm2/rewrite/native_gates.py
diff --git a/src/bloqade/qasm2/rewrite/parallel_to_uop.py b/src/bloqade/qasm2/rewrite/parallel_to_uop.py
diff --git a/src/bloqade/squin/groups.py b/src/bloqade/squin/groups.py
diff --git a/src/bloqade/squin/qubit.py b/src/bloqade/squin/qubit.py

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@`
`10`	`10`	`DeadCodeElimination,`
`11`	`11`	`CommonSubexpressionElimination,`
`12`	`12`	`)`
`13`		`-from kirin.rewrite.result import RewriteResult`
	`13`	`+from kirin.rewrite.abc import RewriteResult`
`14`	`14`
`15`	`15`	`from bloqade.noise import native`
`16`	`16`	`from bloqade.analysis import address`