Refactoring stdlib to fix rewrite

Author: weinbe58

src/bloqade/native/stdlib/broadcast.py

@@ -21,6 +21,21 @@ def _radian_to_turn(angle: float) -> float:
     return angle / (2 * math.pi)
 
 
+@kernel
+def _rx_turns(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    native.r(0.0, angle, qubits)
+
+
+@kernel
+def _ry_turns(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    native.r(0.25, angle, qubits)
+
+
+@kernel
+def _rz_turns(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    native.r(0.5, angle, qubits)
+
+
 @kernel
 def rx(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
     """Apply an RX rotation gate on a group of qubits.
@@ -29,7 +44,7 @@ def rx(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
         angle (float): Rotation angle in radians.
         qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
     """
-    native.r(0.0, _radian_to_turn(angle), qubits)
+    _rx_turns(_radian_to_turn(angle), qubits)
 
 
 @kernel
@@ -70,7 +85,7 @@ def ry(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
         angle (float): Rotation angle in radians.
         qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
     """
-    native.r(0.25, _radian_to_turn(angle), qubits)
+    _ry_turns(_radian_to_turn(angle), qubits)
 
 
 @kernel

src/bloqade/native/upstream/squin2native.py

@@ -20,9 +20,9 @@ class GateRule(RewriteRule):
         stmts.T: (broadcast.t, broadcast.t_adj),
         stmts.SqrtX: (broadcast.sqrt_x, broadcast.sqrt_x_adj),
         stmts.SqrtY: (broadcast.sqrt_y, broadcast.sqrt_y_adj),
-        stmts.Rx: (broadcast.rx,),
-        stmts.Ry: (broadcast.ry,),
-        stmts.Rz: (broadcast.rz,),
+        stmts.Rx: (broadcast._rx_turns,),
+        stmts.Ry: (broadcast._ry_turns,),
+        stmts.Rz: (broadcast._rz_turns,),
         stmts.CX: (broadcast.cx,),
         stmts.CY: (broadcast.cy,),
         stmts.CZ: (broadcast.cz,),

test/native/upstream/test_squin2native.py

@@ -2,9 +2,11 @@
 import pytest
 from kirin.analysis import callgraph
 
-from bloqade import squin
+from bloqade import squin, native
 from bloqade.squin import gate
 from bloqade.pyqrack import StackMemorySimulator
+from bloqade.test_utils import assert_methods
+from bloqade.rewrite.passes import AggressiveUnroll
 from bloqade.native.dialects import gate as native_gate
 from bloqade.native.upstream import GateRule, SquinToNative
 
@@ -48,3 +50,24 @@ def main():
     new_sv /= new_sv[imax := np.abs(new_sv).argmax()] / np.abs(new_sv[imax])
 
     assert np.allclose(old_sv, new_sv)
+
+
+def test_pipeline():
+
+    @squin.kernel
+    def ghz(angle: float):
+        qubits = squin.qalloc(1)
+        squin.rz(angle, qubits[0])
+
+    @native.kernel
+    def ghz_native(angle: float):
+        qubits = squin.qalloc(1)
+        native.rz(angle, qubits[0])
+
+    ghz_native_rewrite = SquinToNative().emit(ghz)
+    AggressiveUnroll(ghz.dialects).fixpoint(ghz_native_rewrite)
+    ghz_native_rewrite.print()
+
+    AggressiveUnroll(ghz_native.dialects).fixpoint(ghz_native)
+    ghz_native.print()
+    assert_methods(ghz_native_rewrite, ghz_native)