Merge origin/david/571-kirin-upgrade-branch into dl/codegen

Author: zhenrongliew

.github/workflows/ci.yml

@@ -43,7 +43,7 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }} # required
       - name: Archive code coverage results
         if: matrix.python-version == '3.12'
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v5
         with:
           name: code-coverage-report
           path: coverage.xml
@@ -55,7 +55,7 @@ jobs:
     steps:
         - uses: actions/checkout@v5
         - name: download covearge
-          uses: actions/download-artifact@v5
+          uses: actions/download-artifact@v6
           with:
             name: code-coverage-report
         - name: check coverage

.github/workflows/release.yml

@@ -19,7 +19,7 @@ jobs:
     - name: Build distribution 📦
       run: uv build
     - name: Store the distribution packages
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       with:
         name: python-package-distributions
         path: dist/
@@ -39,7 +39,7 @@ jobs:
 
     steps:
     - name: Download all the dists
-      uses: actions/download-artifact@v5
+      uses: actions/download-artifact@v6
       with:
         name: python-package-distributions
         path: dist/
@@ -60,12 +60,12 @@ jobs:
 
     steps:
     - name: Download all the dists
-      uses: actions/download-artifact@v5
+      uses: actions/download-artifact@v6
       with:
         name: python-package-distributions
         path: dist/
     - name: Sign the dists with Sigstore
-      uses: sigstore/gh-action-sigstore-python@v3.0.1
+      uses: sigstore/gh-action-sigstore-python@v3.1.0
       with:
         inputs: >-
           ./dist/*.tar.gz

src/bloqade/cirq_utils/emit/base.py

@@ -172,9 +172,6 @@ class EmitCirq(EmitABC[EmitCirqFrame, cirq.Circuit]):
     dialects: ir.DialectGroup = field(default_factory=_default_kernel)
     void = cirq.Circuit()
     qubits: Sequence[cirq.Qid] | None = None
-    _cached_invokes: dict[int, cirq.FrozenCircuit] = field(
-        init=False, default_factory=dict
-    )
 
     def initialize(self) -> Self:
         return super().initialize()

src/bloqade/native/upstream/squin2native.py

@@ -1,14 +1,13 @@
 from itertools import chain
-from dataclasses import field, dataclass
 
-from kirin import ir, passes, rewrite
+from kirin import ir, rewrite
 from kirin.dialects import py, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
-from kirin.passes.callgraph import CallGraphPass, ReplaceMethods
 from kirin.analysis.callgraph import CallGraph
 
 from bloqade.native import kernel, broadcast
 from bloqade.squin.gate import stmts, dialect as gate_dialect
+from bloqade.rewrite.passes import CallGraphPass, UpdateDialectsOnCallGraph
 
 
 class GateRule(RewriteRule):
@@ -46,63 +45,6 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         return RewriteResult(has_done_something=True)
 
 
-@dataclass
-class UpdateDialectsOnCallGraph(passes.Pass):
-    """Update All dialects on the call graph to a new set of dialects given to this pass.
-
-    Usage:
-        pass_ = UpdateDialectsOnCallGraph(rule=rule, dialects=new_dialects)
-        pass_(some_method)
-
-    Note: This pass does not update the dialects of the input method, but copies
-    all other methods invoked within it before updating their dialects.
-
-    """
-
-    fold_pass: passes.Fold = field(init=False)
-
-    def __post_init__(self):
-        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
-
-    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
-        mt_map = {}
-
-        cg = CallGraph(mt)
-
-        all_methods = set(sum(map(tuple, cg.defs.values()), ()))
-        for original_mt in all_methods:
-            if original_mt is mt:
-                new_mt = original_mt
-            else:
-                new_mt = original_mt.similar(self.dialects)
-            mt_map[original_mt] = new_mt
-
-        result = RewriteResult()
-
-        for _, new_mt in mt_map.items():
-            result = (
-                rewrite.Walk(ReplaceMethods(mt_map)).rewrite(new_mt.code).join(result)
-            )
-            self.fold_pass(new_mt)
-
-        return result
-
-
-@dataclass
-class SquinToNativePass(passes.Pass):
-
-    call_graph_pass: CallGraphPass = field(init=False)
-
-    def __post_init__(self):
-        rule = rewrite.Walk(GateRule())
-        self.call_graph_pass = CallGraphPass(
-            self.dialects, rule, no_raise=self.no_raise
-        )
-
-    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
-        return self.call_graph_pass.unsafe_run(mt)
-
-
 class SquinToNative:
     """A Target that converts Squin gates to native gates."""
 
@@ -126,11 +68,10 @@ def emit(self, mt: ir.Method, *, no_raise=True) -> ir.Method:
 
         out = mt.similar(new_dialects)
         UpdateDialectsOnCallGraph(new_dialects, no_raise=no_raise)(out)
-        SquinToNativePass(new_dialects, no_raise=no_raise)(out)
+        CallGraphPass(new_dialects, rewrite.Walk(GateRule()), no_raise=no_raise)(out)
         # verify all kernels in the callgraph
         new_callgraph = CallGraph(out)
-        all_kernels = (ker for kers in new_callgraph.defs.values() for ker in kers)
-        for ker in all_kernels:
+        for ker in new_callgraph.edges.keys():
             ker.verify()
 
         return out

src/bloqade/rewrite/passes/__init__.py

@@ -1,2 +1,7 @@
+from .callgraph import (
+    CallGraphPass as CallGraphPass,
+    ReplaceMethods as ReplaceMethods,
+    UpdateDialectsOnCallGraph as UpdateDialectsOnCallGraph,
+)
 from .aggressive_unroll import AggressiveUnroll as AggressiveUnroll
 from .canonicalize_ilist import CanonicalizeIList as CanonicalizeIList

src/bloqade/rewrite/passes/callgraph.py

@@ -0,0 +1,116 @@
+from dataclasses import field, dataclass
+
+from kirin import ir, passes, rewrite
+from kirin.analysis import CallGraph
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.func.stmts import Invoke
+
+
+@dataclass
+class ReplaceMethods(RewriteRule):
+    new_symbols: dict[ir.Method, ir.Method]
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if (
+            not isinstance(node, Invoke)
+            or (new_callee := self.new_symbols.get(node.callee)) is None
+        ):
+            return RewriteResult()
+
+        node.replace_by(
+            Invoke(
+                inputs=node.inputs,
+                callee=new_callee,
+                purity=node.purity,
+            )
+        )
+
+        return RewriteResult(has_done_something=True)
+
+
+@dataclass
+class UpdateDialectsOnCallGraph(passes.Pass):
+    """Update All dialects on the call graph to a new set of dialects given to this pass.
+
+    Usage:
+        pass_ = UpdateDialectsOnCallGraph(rule=rule, dialects=new_dialects)
+        pass_(some_method)
+
+    Note: This pass does not update the dialects of the input method, but copies
+    all other methods invoked within it before updating their dialects.
+
+    """
+
+    fold_pass: passes.Fold = field(init=False)
+
+    def __post_init__(self):
+        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        mt_map = {}
+
+        cg = CallGraph(mt)
+
+        all_methods = set(sum(map(tuple, cg.defs.values()), ()))
+        for original_mt in all_methods:
+            if original_mt is mt:
+                new_mt = original_mt
+            else:
+                new_mt = original_mt.similar(self.dialects)
+            mt_map[original_mt] = new_mt
+
+        result = RewriteResult()
+
+        for _, new_mt in mt_map.items():
+            result = (
+                rewrite.Walk(ReplaceMethods(mt_map)).rewrite(new_mt.code).join(result)
+            )
+            self.fold_pass(new_mt)
+
+        return result
+
+
+@dataclass
+class CallGraphPass(passes.Pass):
+    """Copy all functions in the call graph and apply a rule to each of them.
+
+
+    Usage:
+        rule = Walk(SomeRewriteRule())
+        pass_ = CallGraphPass(rule=rule, dialects=...)
+        pass_(some_method)
+
+    Note: This pass modifies the input method in place, but copies
+    all methods invoked within it before applying the rule to them.
+
+    """
+
+    rule: RewriteRule
+    """The rule to apply to each function in the call graph."""
+
+    fold_pass: passes.Fold = field(init=False)
+
+    def __post_init__(self):
+        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        result = RewriteResult()
+        mt_map = {}
+
+        cg = CallGraph(mt)
+
+        all_methods = set(cg.edges.keys())
+        for original_mt in all_methods:
+            if original_mt is mt:
+                new_mt = original_mt
+            else:
+                new_mt = original_mt.similar()
+            result = self.rule.rewrite(new_mt.code).join(result)
+            mt_map[original_mt] = new_mt
+
+        if result.has_done_something:
+            for _, new_mt in mt_map.items():
+                rewrite.Walk(ReplaceMethods(mt_map)).rewrite(new_mt.code)
+                self.fold_pass(new_mt)
+
+        return result

src/bloqade/squin/rewrite/U3_to_clifford.py

@@ -15,6 +15,14 @@ class Sdag(ir.Statement):
     pass
 
 
+class SqrtXdag(ir.Statement):
+    pass
+
+
+class SqrtYdag(ir.Statement):
+    pass
+
+
 # (theta, phi, lam)
 U3_HALF_PI_ANGLE_TO_GATES: dict[
     tuple[int, int, int], list[type[ir.Statement]] | list[None]
@@ -27,21 +35,21 @@ class Sdag(ir.Statement):
     (1, 0, 1): [gate.stmts.S, gate.stmts.SqrtY],
     (1, 0, 2): [gate.stmts.H],
     (1, 0, 3): [Sdag, gate.stmts.SqrtY],
-    (1, 1, 0): [gate.stmts.SqrtY, gate.stmts.S],
-    (1, 1, 1): [gate.stmts.S, gate.stmts.SqrtY, gate.stmts.S],
-    (1, 1, 2): [gate.stmts.Z, gate.stmts.SqrtY, gate.stmts.S],
-    (1, 1, 3): [Sdag, gate.stmts.SqrtY, gate.stmts.S],
-    (1, 2, 0): [gate.stmts.SqrtY, gate.stmts.Z],
-    (1, 2, 1): [gate.stmts.S, gate.stmts.SqrtY, gate.stmts.Z],
-    (1, 2, 2): [gate.stmts.Z, gate.stmts.SqrtY, gate.stmts.Z],
-    (1, 2, 3): [Sdag, gate.stmts.SqrtY, gate.stmts.Z],
-    (1, 3, 0): [gate.stmts.SqrtY, Sdag],
-    (1, 3, 1): [gate.stmts.S, gate.stmts.SqrtY, Sdag],
-    (1, 3, 2): [gate.stmts.Z, gate.stmts.SqrtY, Sdag],
-    (1, 3, 3): [Sdag, gate.stmts.SqrtY, Sdag],
+    (1, 1, 0): [gate.stmts.S, SqrtXdag],
+    (1, 1, 1): [gate.stmts.Z, SqrtXdag],
+    (1, 1, 2): [Sdag, SqrtXdag],
+    (1, 1, 3): [SqrtXdag],
+    (1, 2, 0): [gate.stmts.Z, SqrtYdag],
+    (1, 2, 1): [Sdag, SqrtYdag],
+    (1, 2, 2): [SqrtYdag],
+    (1, 2, 3): [gate.stmts.S, SqrtYdag],
+    (1, 3, 0): [Sdag, gate.stmts.SqrtX],
+    (1, 3, 1): [gate.stmts.SqrtX],
+    (1, 3, 2): [gate.stmts.S, gate.stmts.SqrtX],
+    (1, 3, 3): [gate.stmts.Z, gate.stmts.SqrtX],
     (2, 0, 0): [gate.stmts.Y],
     (2, 0, 1): [gate.stmts.S, gate.stmts.Y],
-    (2, 0, 2): [gate.stmts.Z, gate.stmts.Y],
+    (2, 0, 2): [gate.stmts.X],
     (2, 0, 3): [Sdag, gate.stmts.Y],
 }
 
@@ -106,6 +114,10 @@ def rewrite_U3(self, node: gate.stmts.U3) -> RewriteResult:
         for gate_stmt in gates:
             if gate_stmt is Sdag:
                 new_stmt = gate.stmts.S(adjoint=True, qubits=node.qubits)
+            elif gate_stmt is SqrtXdag:
+                new_stmt = gate.stmts.SqrtX(adjoint=True, qubits=node.qubits)
+            elif gate_stmt is SqrtYdag:
+                new_stmt = gate.stmts.SqrtY(adjoint=True, qubits=node.qubits)
             else:
                 new_stmt = gate_stmt(qubits=node.qubits)
             new_stmt.insert_before(node)

src/bloqade/stim/emit/__init__.py

@@ -1 +1,2 @@
+from . import impls as impls
 from .stim_str import FuncEmit as FuncEmit, EmitStimMain as EmitStimMain

src/bloqade/stim/emit/impls.py

@@ -0,0 +1,17 @@
+from kirin.emit import EmitStrFrame
+from kirin.interp import MethodTable, impl
+from kirin.dialects.debug import Info, dialect
+
+from bloqade.stim.emit.stim_str import EmitStimMain
+
+
+@dialect.register(key="emit.stim")
+class EmitStimDebugMethods(MethodTable):
+
+    @impl(Info)
+    def info(self, emit: EmitStimMain, frame: EmitStrFrame, stmt: Info):
+
+        msg: str = frame.get(stmt.msg)
+        emit.writeln(frame, f"# {msg}")
+
+        return ()

test/cirq_utils/test_cirq_to_squin.py

@@ -402,15 +402,6 @@ def multi_arg(n: int, p: float):
     print(circuit)
 
 
-if __name__ == "__main__":
-    test_kernel_with_args()
-
-
-@pytest.mark.xfail
-def test_amplitude_damping():
-    test_circuit(amplitude_damping)
-
-
 def test_trotter():
 
     # NOTE: stolen from jonathan's tutorial