Adding Initialize statement for logical qubits.

src/bloqade/gemini/dialects/logical/stmts.py

@@ -6,6 +6,26 @@
 
 from ._dialect import dialect
 
+
+@statement(dialect=dialect)
+class Initialize(ir.Statement):
+    """Initialize a list of logical qubits to an arbitrary state.
+
+    Args:
+        phi (float): Angle for rotation around the Z axis
+        theta: float angle for rotation around the Y axis
+        phi: float angle for rotation around the Z axis
+        qubits (IList[QubitType, Len]): The list of logical qubits to initialize
+
+    """
+
+    traits = frozenset({})
+    theta: ir.SSAValue = info.argument(types.Float)
+    phi: ir.SSAValue = info.argument(types.Float)
+    lam: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
 Len = types.TypeVar("Len")
 CodeN = types.TypeVar("CodeN")
 

src/bloqade/gemini/rewrite/initialize.py

@@ -0,0 +1,30 @@
+from kirin import ir
+from kirin.rewrite import abc as rewrite_abc
+
+from bloqade.squin.gate.stmts import U3
+
+from ..dialects.logical.stmts import Initialize
+
+
+class __RewriteU3ToInitialize(rewrite_abc.RewriteRule):
+    """Rewrite U3 gates to Initialize statements.
+
+    Note:
+
+    This rewrite is only valid in the context of logical qubits, where the U3 gate
+    can be interpreted as initializing a qubit to an arbitrary state.
+
+    The U3 gate with parameters (theta, phi, lam) is equivalent to initializing
+    a qubit to the state defined by those angles.
+
+    This rewrite also assumes there are no other U3 gates acting on the same qubits
+    later in the circuit, as that would conflict with the initialization semantics.
+
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+        if not isinstance(node, U3):
+            return rewrite_abc.RewriteResult()
+
+        node.replace_by(Initialize(*node.args))
+        return rewrite_abc.RewriteResult(has_done_something=True)

test/gemini/test_rewrite.py

@@ -0,0 +1,30 @@
+from kirin import ir, rewrite
+from kirin.dialects import py
+
+from bloqade.test_utils import assert_nodes
+from bloqade.squin.gate.stmts import U3
+from bloqade.gemini.rewrite.initialize import __RewriteU3ToInitialize
+from bloqade.gemini.dialects.logical.stmts import Initialize
+
+
+def test_rewrite_u3_to_initialize():
+    theta = ir.TestValue()
+    phi = ir.TestValue()
+    qubits = ir.TestValue()
+    test_block = ir.Block(
+        [
+            lam_stmt := py.Constant(1.0),
+            U3(theta, phi, lam_stmt.result, qubits),
+        ]
+    )
+
+    expected_block = ir.Block(
+        [
+            lam_stmt := py.Constant(1.0),
+            Initialize(theta, phi, lam_stmt.result, qubits),
+        ]
+    )
+
+    rewrite.Walk(__RewriteU3ToInitialize()).rewrite(test_block)
+
+    assert_nodes(test_block, expected_block)