merging main

Author: weinbe58

src/bloqade/cirq_utils/emit/gate.py

@@ -86,3 +86,19 @@ def rot(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.RotationGat
 
         frame.circuit.append(cirq_op.on_each(qubits))
         return ()
+
+    @impl(gate.stmts.U3)
+    def u3(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: gate.stmts.U3):
+        qubits = frame.get(stmt.qubits)
+
+        theta = frame.get(stmt.theta) * 2 * math.pi
+        phi = frame.get(stmt.phi) * 2 * math.pi
+        lam = frame.get(stmt.lam) * 2 * math.pi
+
+        frame.circuit.append(cirq.Rz(rads=lam).on_each(*qubits))
+
+        frame.circuit.append(cirq.Ry(rads=theta).on_each(*qubits))
+
+        frame.circuit.append(cirq.Rz(rads=phi).on_each(*qubits))
+
+        return ()

src/bloqade/native/__init__.py

@@ -14,7 +14,6 @@
     ry as ry,
     rz as rz,
     u3 as u3,
-    rot as rot,
     s_dag as s_dag,
     shift as shift,
     sqrt_x as sqrt_x,

src/bloqade/native/stdlib/broadcast.py

@@ -187,33 +187,23 @@ def shift(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
     rz(angle / 2.0, qubits)
 
 
-@kernel
-def rot(phi: float, theta: float, omega: float, qubits: ilist.IList[qubit.Qubit, Any]):
-    """Apply a general single-qubit rotation on a group of qubits.
-
-    Args:
-        phi (float): Z rotation before Y (radians).
-        theta (float): Y rotation (radians).
-        omega (float): Z rotation after Y (radians).
-        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
-    """
-    rz(phi, qubits)
-    ry(theta, qubits)
-    rz(omega, qubits)
-
-
 @kernel
 def u3(theta: float, phi: float, lam: float, qubits: ilist.IList[qubit.Qubit, Any]):
     """Apply the U3 gate on a group of qubits.
 
+    The applied gate is represented by the unitary matrix given by:
+
+    $$ U3(\\theta, \\phi, \\lambda) = R_z(\\phi)R_y(\\theta)R_z(\\lambda) $$
+
     Args:
         theta (float): Rotation around Y axis (radians).
         phi (float): Global phase shift component (radians).
         lam (float): Z rotations in decomposition (radians).
         qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
     """
-    rot(lam, theta, -lam, qubits)
-    shift(phi + lam, qubits)
+    rz(lam, qubits)
+    ry(theta, qubits)
+    rz(phi, qubits)
 
 
 N = TypeVar("N")

src/bloqade/native/stdlib/simple.py

@@ -170,23 +170,14 @@ def shift(angle: float, qubit: qubit.Qubit):
     broadcast.shift(angle, ilist.IList([qubit]))
 
 
-@kernel
-def rot(phi: float, theta: float, omega: float, qubit: qubit.Qubit):
-    """Apply a general single-qubit rotation on a single qubit.
-
-    Args:
-        phi (float): Z rotation before Y (radians).
-        theta (float): Y rotation (radians).
-        omega (float): Z rotation after Y (radians).
-        qubit (qubit.Qubit): The qubit to apply the rotation to.
-    """
-    broadcast.rot(phi, theta, omega, ilist.IList([qubit]))
-
-
 @kernel
 def u3(theta: float, phi: float, lam: float, qubit: qubit.Qubit):
     """Apply the U3 gate on a single qubit.
 
+    The applied gate is represented by the unitary matrix given by:
+
+    $$ U3(\\theta, \\phi, \\lambda) = R_z(\\phi)R_y(\\theta)R_z(\\lambda) $$
+
     Args:
         theta (float): Rotation angle around the Y axis in radians.
         phi (float): Rotation angle around the Z axis in radians.

src/bloqade/native/upstream/__init__.py

@@ -1,4 +1,5 @@
 from .squin2native import (
+    GateRule as GateRule,
     SquinToNative as SquinToNative,
     SquinToNativePass as SquinToNativePass,
 )

src/bloqade/native/upstream/squin2native.py

@@ -27,6 +27,7 @@ class GateRule(RewriteRule):
         stmts.CX: (broadcast.cx,),
         stmts.CY: (broadcast.cy,),
         stmts.CZ: (broadcast.cz,),
+        stmts.U3: (broadcast.u3,),
     }
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:

src/bloqade/pyqrack/squin/gate/gate.py

@@ -12,6 +12,7 @@
     CX,
     CY,
     CZ,
+    U3,
     H,
     S,
     T,
@@ -120,3 +121,16 @@ def control(
         for control, target in zip(controls, targets):
             if control.is_active() and target.is_active():
                 getattr(control.sim_reg, method_name)([control.addr], target.addr)
+
+    @interp.impl(U3)
+    def u3(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: U3):
+        theta = frame.get(stmt.theta) * 2 * math.pi
+        phi = frame.get(stmt.phi) * 2 * math.pi
+        lam = frame.get(stmt.lam) * 2 * math.pi
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+
+        for qbit in qubits:
+            if not qbit.is_active():
+                continue
+
+            qbit.sim_reg.u(qbit.addr, theta, phi, lam)

src/bloqade/qasm2/passes/fold.py

@@ -1,27 +1,12 @@
 from dataclasses import field, dataclass
 
 from kirin import ir
-from kirin.passes import Pass, TypeInfer
-from kirin.rewrite import (
-    Walk,
-    Chain,
-    Inline,
-    Fixpoint,
-    WrapConst,
-    Call2Invoke,
-    ConstantFold,
-    CFGCompactify,
-    InlineGetItem,
-    InlineGetField,
-    DeadCodeElimination,
-    CommonSubexpressionElimination,
-)
-from kirin.analysis import const
-from kirin.dialects import scf, ilist
+from kirin.passes import Pass
 from kirin.ir.method import Method
 from kirin.rewrite.abc import RewriteResult
 
 from bloqade.qasm2.dialects import expr
+from bloqade.rewrite.passes import AggressiveUnroll
 
 from .unroll_if import UnrollIfs
 
@@ -30,71 +15,27 @@
 class QASM2Fold(Pass):
     """Fold pass for qasm2.extended"""
 
-    constprop: const.Propagate = field(init=False)
     inline_gate_subroutine: bool = True
     unroll_ifs: bool = True
+    aggressive_unroll: AggressiveUnroll = field(init=False)
 
     def __post_init__(self):
-        self.constprop = const.Propagate(self.dialects)
-        self.typeinfer = TypeInfer(self.dialects)
+        def inline_simple(node: ir.Statement):
+            if isinstance(node, expr.GateFunction):
+                return self.inline_gate_subroutine
 
-    def unsafe_run(self, mt: Method) -> RewriteResult:
-        result = RewriteResult()
-        frame, _ = self.constprop.run_analysis(mt)
-        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
-        rule = Chain(
-            ConstantFold(),
-            Call2Invoke(),
-            InlineGetField(),
-            InlineGetItem(),
-            DeadCodeElimination(),
-            CommonSubexpressionElimination(),
-        )
-        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+            return True
 
-        result = (
-            Walk(
-                Chain(
-                    scf.unroll.PickIfElse(),
-                    scf.unroll.ForLoop(),
-                    scf.trim.UnusedYield(),
-                )
-            )
-            .rewrite(mt.code)
-            .join(result)
+        self.aggressive_unroll = AggressiveUnroll(
+            self.dialects, inline_simple, no_raise=self.no_raise
         )
 
-        if self.unroll_ifs:
-            UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
-
-        # run typeinfer again after unroll etc. because we now insert
-        # a lot of new nodes, which might have more precise types
-        self.typeinfer.unsafe_run(mt)
-        result = (
-            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
-            .rewrite(mt.code)
-            .join(result)
-        )
-
-        def inline_simple(node: ir.Statement):
-            if isinstance(node, expr.GateFunction):
-                return self.inline_gate_subroutine
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
 
-            if not isinstance(node.parent_stmt, (scf.For, scf.IfElse)):
-                return True  # always inline calls outside of loops and if-else
+        if self.unroll_ifs:
+            result = UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
 
-            # inside loops and if-else, only inline simple functions, i.e. functions with a single block
-            if (trait := node.get_trait(ir.CallableStmtInterface)) is None:
-                return False  # not a callable, don't inline to be safe
-            region = trait.get_callable_region(node)
-            return len(region.blocks) == 1
+        result = self.aggressive_unroll.unsafe_run(mt).join(result)
 
-        result = (
-            Walk(
-                Inline(inline_simple),
-            )
-            .rewrite(mt.code)
-            .join(result)
-        )
-        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
         return result

src/bloqade/rewrite/passes/__init__.py

@@ -1 +1,2 @@
+from .aggressive_unroll import AggressiveUnroll as AggressiveUnroll
 from .canonicalize_ilist import CanonicalizeIList as CanonicalizeIList

src/bloqade/rewrite/passes/aggressive_unroll.py

@@ -0,0 +1,93 @@
+from typing import Callable
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass, HintConst, TypeInfer
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Inline,
+    Fixpoint,
+    Call2Invoke,
+    ConstantFold,
+    CFGCompactify,
+    InlineGetItem,
+    InlineGetField,
+    DeadCodeElimination,
+)
+from kirin.dialects import scf, ilist
+from kirin.ir.method import Method
+from kirin.rewrite.abc import RewriteResult
+from kirin.rewrite.cse import CommonSubexpressionElimination
+from kirin.passes.aggressive import UnrollScf
+
+
+@dataclass
+class Fold(Pass):
+    hint_const: HintConst = field(init=False)
+
+    def __post_init__(self):
+        self.hint_const = HintConst(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        result = self.hint_const.unsafe_run(mt).join(result)
+        rule = Chain(
+            ConstantFold(),
+            Call2Invoke(),
+            InlineGetField(),
+            InlineGetItem(),
+            ilist.rewrite.InlineGetItem(),
+            ilist.rewrite.HintLen(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+        return result
+
+
+@dataclass
+class AggressiveUnroll(Pass):
+    """A pass to unroll structured control flow"""
+
+    additional_inline_heuristic: Callable[[ir.Statement], bool] = lambda node: True
+
+    fold: Fold = field(init=False)
+    typeinfer: TypeInfer = field(init=False)
+    scf_unroll: UnrollScf = field(init=False)
+
+    def __post_init__(self):
+        self.fold = Fold(self.dialects, no_raise=self.no_raise)
+        self.typeinfer = TypeInfer(self.dialects, no_raise=self.no_raise)
+        self.scf_unroll = UnrollScf(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        result = self.scf_unroll.unsafe_run(mt).join(result)
+        result = (
+            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        result = self.typeinfer.unsafe_run(mt).join(result)
+        result = self.fold.unsafe_run(mt).join(result)
+        result = Walk(Inline(self.inline_heuristic)).rewrite(mt.code).join(result)
+        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
+
+        rule = Chain(
+            CommonSubexpressionElimination(),
+            DeadCodeElimination(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+        return result
+
+    def inline_heuristic(self, node: ir.Statement) -> bool:
+        """The heuristic to decide whether to inline a function call or not.
+        inside loops and if-else, only inline simple functions, i.e.
+        functions with a single block
+        """
+        return not isinstance(
+            node.parent_stmt, (scf.For, scf.IfElse)
+        ) and self.additional_inline_heuristic(
+            node
+        )  # always inline calls outside of loops and if-else