From 25707861a7e6bde56309c01b36ca3d6c33e0d6bb Mon Sep 17 00:00:00 2001 From: Phillip Weinberg Date: Fri, 5 Dec 2025 14:50:25 -0500 Subject: [PATCH 1/4] adding constant statement to explore no-op path --- src/bloqade/gemini/dialects/logical/stmts.py | 20 +++++++++++++ src/bloqade/gemini/rewrite/__init__.py | 0 src/bloqade/gemini/rewrite/initialize.py | 30 ++++++++++++++++++++ test/gemini/test_rewrite.py | 30 ++++++++++++++++++++ 4 files changed, 80 insertions(+) create mode 100644 src/bloqade/gemini/rewrite/__init__.py create mode 100644 src/bloqade/gemini/rewrite/initialize.py create mode 100644 test/gemini/test_rewrite.py diff --git a/src/bloqade/gemini/dialects/logical/stmts.py b/src/bloqade/gemini/dialects/logical/stmts.py index 873c1bfd6..153650337 100644 --- a/src/bloqade/gemini/dialects/logical/stmts.py +++ b/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: + qubits (IList[QubitType, Len]): The list of logical qubits to initialize + 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 + + """ + + 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") diff --git a/src/bloqade/gemini/rewrite/__init__.py b/src/bloqade/gemini/rewrite/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/bloqade/gemini/rewrite/initialize.py b/src/bloqade/gemini/rewrite/initialize.py new file mode 100644 index 000000000..8202bed47 --- /dev/null +++ b/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) diff --git a/test/gemini/test_rewrite.py b/test/gemini/test_rewrite.py new file mode 100644 index 000000000..97082824d --- /dev/null +++ b/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) From c326185cfdaabe2cf4b5d80b062e88a68c2fb15a Mon Sep 17 00:00:00 2001 From: Phillip Weinberg Date: Fri, 5 Dec 2025 14:51:24 -0500 Subject: [PATCH 2/4] fixing doc string --- src/bloqade/gemini/dialects/logical/stmts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/bloqade/gemini/dialects/logical/stmts.py b/src/bloqade/gemini/dialects/logical/stmts.py index 153650337..5694576a4 100644 --- a/src/bloqade/gemini/dialects/logical/stmts.py +++ b/src/bloqade/gemini/dialects/logical/stmts.py @@ -12,10 +12,10 @@ class Initialize(ir.Statement): """Initialize a list of logical qubits to an arbitrary state. Args: - qubits (IList[QubitType, Len]): The list of logical qubits to initialize 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 """ From 2b823f069aae446c19ab70b78f1077bd43ae4976 Mon Sep 17 00:00:00 2001 From: Phillip Weinberg Date: Sat, 6 Dec 2025 11:45:08 -0500 Subject: [PATCH 3/4] fixing issues with terminal measure --- src/bloqade/gemini/dialects/logical/_interface.py | 5 ++--- src/bloqade/gemini/dialects/logical/stmts.py | 3 +-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/bloqade/gemini/dialects/logical/_interface.py b/src/bloqade/gemini/dialects/logical/_interface.py index b9bb692fb..de44390a1 100644 --- a/src/bloqade/gemini/dialects/logical/_interface.py +++ b/src/bloqade/gemini/dialects/logical/_interface.py @@ -1,4 +1,4 @@ -from typing import TypeVar +from typing import Any, TypeVar from kirin import lowering from kirin.dialects import ilist @@ -8,13 +8,12 @@ from .stmts import TerminalLogicalMeasurement Len = TypeVar("Len") -CodeN = TypeVar("CodeN") @lowering.wraps(TerminalLogicalMeasurement) def terminal_measure( qubits: ilist.IList[Qubit, Len], -) -> ilist.IList[ilist.IList[MeasurementResult, CodeN], Len]: +) -> ilist.IList[ilist.IList[MeasurementResult, Any], Len]: """Perform measurements on a list of logical qubits. Measurements are returned as a nested list where each member list diff --git a/src/bloqade/gemini/dialects/logical/stmts.py b/src/bloqade/gemini/dialects/logical/stmts.py index 5694576a4..f26575eb1 100644 --- a/src/bloqade/gemini/dialects/logical/stmts.py +++ b/src/bloqade/gemini/dialects/logical/stmts.py @@ -27,7 +27,6 @@ class Initialize(ir.Statement): Len = types.TypeVar("Len") -CodeN = types.TypeVar("CodeN") @statement(dialect=dialect) @@ -48,5 +47,5 @@ class TerminalLogicalMeasurement(ir.Statement): traits = frozenset({lowering.FromPythonCall()}) qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, Len]) result: ir.ResultValue = info.result( - ilist.IListType[ilist.IListType[MeasurementResultType, CodeN], Len] + ilist.IListType[ilist.IListType[MeasurementResultType, types.Any], Len] ) From 8e3415fb2d394d22f0310597f13e83d619c1aaaa Mon Sep 17 00:00:00 2001 From: Phillip Weinberg Date: Sun, 7 Dec 2025 15:15:13 -0500 Subject: [PATCH 4/4] fixing doc string --- src/bloqade/gemini/dialects/logical/stmts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/bloqade/gemini/dialects/logical/stmts.py b/src/bloqade/gemini/dialects/logical/stmts.py index f26575eb1..f43ef8f09 100644 --- a/src/bloqade/gemini/dialects/logical/stmts.py +++ b/src/bloqade/gemini/dialects/logical/stmts.py @@ -13,8 +13,8 @@ class Initialize(ir.Statement): 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 + 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 """