Rewrite gemini.logicalInitialize to place.Initialize. (#50)

* rename rewrite and flattening qubit ids for CZ

* removing helper function

* Adding merge for initialize

* Renaming test file

* Adding test for initialize

* removing cached property to prevent potential bugs

Author: weinbe58

src/bloqade/lanes/dialects/place.py

@@ -41,8 +41,15 @@ class Initialize(QuantumStmt):
 
 @statement(dialect=dialect)
 class CZ(QuantumStmt):
-    targets: tuple[int, ...] = info.attribute()
-    controls: tuple[int, ...] = info.attribute()
+    qubits: tuple[int, ...] = info.attribute()
+
+    @property
+    def controls(self) -> tuple[int, ...]:
+        return self.qubits[: len(self.qubits) // 2]
+
+    @property
+    def targets(self) -> tuple[int, ...]:
+        return self.qubits[len(self.qubits) // 2 :]
 
 
 @statement(dialect=dialect)

src/bloqade/lanes/rewrite/native2place.py renamed to src/bloqade/lanes/rewrite/circuit2place.py

@@ -18,12 +18,20 @@ class RewritePlaceOperations(abc.RewriteRule):
     This is a placeholder for the actual implementation.
     """
 
-    def default_(self, node: ir.Statement) -> abc.RewriteResult:
-        return abc.RewriteResult()
-
     def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
+        if not isinstance(
+            node,
+            (
+                gemini_stmts.TerminalLogicalMeasurement,
+                gemini_stmts.Initialize,
+                gate.CZ,
+                gate.R,
+                gate.Rz,
+            ),
+        ):
+            return abc.RewriteResult()
         rewrite_method_name = f"rewrite_{type(node).__name__}"
-        rewrite_method = getattr(self, rewrite_method_name, self.default_)
+        rewrite_method = getattr(self, rewrite_method_name)
         return rewrite_method(node)
 
     def prep_region(self) -> tuple[ir.Region, ir.Block, ir.SSAValue]:
@@ -46,6 +54,25 @@ def construct_execute(
 
         return place.StaticPlacement(qubits=qubits, body=body)
 
+    def rewrite_Initialize(self, node: gemini_stmts.Initialize) -> abc.RewriteResult:
+        if not isinstance(args_list := node.qubits.owner, ilist.New):
+            return abc.RewriteResult()
+
+        inputs = args_list.values
+        body, block, entry_state = self.prep_region()
+        gate_stmt = place.Initialize(
+            entry_state,
+            phi=node.phi,
+            theta=node.theta,
+            lam=node.lam,
+            qubits=tuple(range(len(inputs))),
+        )
+        node.replace_by(
+            self.construct_execute(gate_stmt, qubits=inputs, body=body, block=block)
+        )
+
+        return abc.RewriteResult(has_done_something=True)
+
     def rewrite_TerminalLogicalMeasurement(
         self, node: gemini_stmts.TerminalLogicalMeasurement
     ) -> abc.RewriteResult:
@@ -82,13 +109,11 @@ def rewrite_CZ(self, node: gate.CZ) -> abc.RewriteResult:
             return abc.RewriteResult()
 
         all_qubits = tuple(range(len(targets) + len(controls)))
-        n_controls = len(controls)
 
         body, block, entry_state = self.prep_region()
         stmt = place.CZ(
             entry_state,
-            controls=all_qubits[:n_controls],
-            targets=all_qubits[n_controls:],
+            qubits=all_qubits,
         )
 
         node.replace_by(
@@ -155,27 +180,6 @@ class MergePlacementRegions(abc.RewriteRule):
     merge_heuristic: Callable[[ir.Region, ir.Region], bool] = _default_merge_heuristic
     """Heuristic function to decide whether to merge two circuit regions."""
 
-    def remap_qubits(
-        self,
-        curr_state: ir.SSAValue,
-        node: place.R | place.Rz | place.CZ | place.EndMeasure,
-        input_map: dict[int, int],
-    ):
-        if isinstance(node, place.CZ):
-            return place.CZ(
-                curr_state,
-                targets=tuple(input_map[i] for i in node.targets),
-                controls=tuple(input_map[i] for i in node.controls),
-            )
-        else:
-            return node.from_stmt(
-                node,
-                args=(curr_state, *node.args[1:]),
-                attributes={
-                    "qubits": ir.PyAttr(tuple(input_map[i] for i in node.qubits))
-                },
-            )
-
     def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if not (
             isinstance(node, place.StaticPlacement)
@@ -206,8 +210,18 @@ def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         curr_yield.delete()
 
         for stmt in next_node.body.blocks[0].stmts:
-            if isinstance(stmt, (place.R, place.Rz, place.CZ, place.EndMeasure)):
-                remapped_stmt = self.remap_qubits(curr_state, stmt, new_input_map)
+            if isinstance(
+                stmt, (place.R, place.Rz, place.CZ, place.EndMeasure, place.Initialize)
+            ):
+                remapped_stmt = stmt.from_stmt(
+                    stmt,
+                    args=(curr_state, *stmt.args[1:]),
+                    attributes={
+                        "qubits": ir.PyAttr(
+                            tuple(new_input_map[i] for i in stmt.qubits)
+                        )
+                    },
+                )
                 curr_state = remapped_stmt.results[0]
                 new_block.stmts.append(remapped_stmt)
                 for old_result, new_result in zip(

src/bloqade/lanes/upstream.py

@@ -10,7 +10,7 @@
 from bloqade.lanes.analysis import layout, placement
 from bloqade.lanes.dialects import move, place
 from bloqade.lanes.passes.canonicalize import CanonicalizeNative
-from bloqade.lanes.rewrite import native2place, place2move
+from bloqade.lanes.rewrite import circuit2place, place2move
 
 
 def default_merge_heuristic(region_a: ir.Region, region_b: ir.Region) -> bool:
@@ -29,11 +29,11 @@ def emit(self, mt: Method, no_raise: bool = True):
         AggressiveUnroll(out.dialects, no_raise=no_raise).fixpoint(out)
         CanonicalizeNative(out.dialects, no_raise=no_raise).fixpoint(out)
         rewrite.Walk(
-            native2place.RewritePlaceOperations(),
+            circuit2place.RewritePlaceOperations(),
         ).rewrite(out.code)
 
         rewrite.Fixpoint(
-            rewrite.Walk(native2place.MergePlacementRegions(self.merge_heuristic))
+            rewrite.Walk(circuit2place.MergePlacementRegions(self.merge_heuristic))
         ).rewrite(out.code)
         passes.TypeInfer(out.dialects)(out)
 

test/rewrite/test_native2circuit.py renamed to test/rewrite/test_circuit2place.py

@@ -6,7 +6,7 @@
 
 from bloqade.lanes import types
 from bloqade.lanes.dialects import place
-from bloqade.lanes.rewrite.native2place import (
+from bloqade.lanes.rewrite.circuit2place import (
     MergePlacementRegions,
     RewritePlaceOperations,
 )
@@ -33,9 +33,7 @@ def test_cz():
     )
 
     entry_state = block.args.append_from(types.StateType, name="entry_state")
-    block.stmts.append(
-        gate_stmt := place.CZ(entry_state, controls=(0, 1), targets=(2, 3))
-    )
+    block.stmts.append(gate_stmt := place.CZ(entry_state, qubits=(0, 1, 2, 3)))
     block.stmts.append(place.Yield(gate_stmt.state_after))
 
     rule = rewrite.Walk(RewritePlaceOperations())
@@ -153,6 +151,53 @@ def test_measurement():
     assert_nodes(test_block, expected_block)
 
 
+def test_initialize():
+    test_block = ir.Block(
+        [
+            qubits := ilist.New(
+                values=(
+                    q0 := ir.TestValue(),
+                    q1 := ir.TestValue(),
+                    q2 := ir.TestValue(),
+                )
+            ),
+            gemini_stmts.Initialize(
+                theta := ir.TestValue(),
+                phi := ir.TestValue(),
+                lam := ir.TestValue(),
+                qubits=qubits.result,
+            ),
+        ],
+    )
+
+    expected_block = ir.Block(
+        [
+            qubits := ilist.New(values=(q0, q1, q2)),
+        ]
+    )
+
+    block = ir.Block()
+
+    entry_state = block.args.append_from(types.StateType, name="entry_state")
+    block.stmts.append(
+        gate_stmt := place.Initialize(
+            entry_state,
+            theta=theta,
+            phi=phi,
+            lam=lam,
+            qubits=(0, 1, 2),
+        )
+    )
+    block.stmts.append(place.Yield(gate_stmt.state_after))
+    expected_block.stmts.append(
+        place.StaticPlacement(qubits=(q0, q1, q2), body=ir.Region(block))
+    )
+    rule = rewrite.Walk(RewritePlaceOperations())
+
+    rule.rewrite(test_block)
+    assert_nodes(test_block, expected_block)
+
+
 def test_merge_regions():
 
     qubits = tuple(ir.TestValue() for _ in range(10))