Adding initial angles for initialize to layout analysis (#47)

* Adding initial angles for initialize to layout analysis

* remove exception

* removing args

Author: weinbe58

src/bloqade/lanes/analysis/layout/analysis.py

@@ -5,7 +5,6 @@
 from kirin import ir
 from kirin.analysis.forward import Forward, ForwardFrame
 from kirin.lattice import EmptyLattice
-from typing_extensions import Self
 
 from bloqade.lanes.layout.encoding import LocationAddress
 
@@ -39,12 +38,15 @@ def compute_layout(
 
 @dataclass
 class LayoutAnalysis(Forward):
-    keys = ("circuit.layout",)
+    keys = ("place.layout",)
     lattice = EmptyLattice
 
     heuristic: LayoutHeuristicABC
     address_entries: dict[ir.SSAValue, address.Address]
     all_qubits: tuple[int, ...] = field(init=False)
+    thetas: dict[int, ir.SSAValue] = field(default_factory=dict, init=False)
+    phis: dict[int, ir.SSAValue] = field(default_factory=dict, init=False)
+    lams: dict[int, ir.SSAValue] = field(default_factory=dict, init=False)
     stages: list[tuple[tuple[int, int], ...]] = field(default_factory=list, init=False)
     global_address_stack: list[int] = field(default_factory=list, init=False)
 
@@ -60,9 +62,12 @@ def __post_init__(self) -> None:
         )
         super().__post_init__()
 
-    def initialize(self) -> Self:
+    def initialize(self):
         self.stages.clear()
         self.global_address_stack.clear()
+        self.thetas.clear()
+        self.phis.clear()
+        self.lams.clear()
         return super().initialize()
 
     def eval_stmt_fallback(self, frame, stmt):
@@ -76,15 +81,36 @@ def add_stage(self, control: tuple[int, ...], target: tuple[int, ...]):
     def method_self(self, method: ir.Method):
         return EmptyLattice.bottom()
 
-    def get_layout_no_raise(self, method: ir.Method) -> tuple[LocationAddress, ...]:
+    def process_results(self):
+        layout = self.heuristic.compute_layout(self.all_qubits, self.stages)
+        init_locations = tuple(
+            loc
+            for qubit_id, loc in zip(self.all_qubits, layout)
+            if qubit_id in self.thetas
+        )
+        thetas = tuple(
+            self.thetas[qubit_id]
+            for qubit_id in self.all_qubits
+            if qubit_id in self.thetas
+        )
+        phis = tuple(
+            self.phis[qubit_id] for qubit_id in self.all_qubits if qubit_id in self.phis
+        )
+        lams = tuple(
+            self.lams[qubit_id] for qubit_id in self.all_qubits if qubit_id in self.lams
+        )
+
+        return layout, init_locations, thetas, phis, lams
+
+    def get_layout_no_raise(self, method: ir.Method):
         """Get the layout for a given method."""
         self.run_no_raise(method)
-        return self.heuristic.compute_layout(self.all_qubits, self.stages)
+        return self.process_results()
 
-    def get_layout(self, method: ir.Method) -> tuple[LocationAddress, ...]:
+    def get_layout(self, method: ir.Method):
         """Get the layout for a given method."""
         self.run(method)
-        return self.heuristic.compute_layout(self.all_qubits, self.stages)
+        return self.process_results()
 
     def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
         return tuple(EmptyLattice.bottom() for _ in node.results)

src/bloqade/lanes/arch/gemini/logical.py

@@ -21,6 +21,28 @@ def get_arch_spec():
 dialect = ir.Dialect("gemini.logical")
 
 
+@decl.statement(dialect=dialect)
+class Fill(ir.Statement):
+
+    filled: ir.SSAValue = info.argument(
+        ilist.IListType[types.Tuple[types.Float, types.Float], types.Any]
+    )
+    vacant: ir.ResultValue = info.result(
+        ilist.IListType[types.Tuple[types.Float, types.Float], types.Any]
+    )
+
+
+NumRows = types.TypeVar("NumRows")
+
+
+@decl.statement(dialect=dialect)
+class LogicalInitialize(ir.Statement):
+    y_locs: ir.SSAValue = info.argument(ilist.IListType[types.Float, NumRows])
+    x_locs: ir.SSAValue = info.argument(
+        ilist.IListType[ilist.IListType[types.Float, types.Any], NumRows]
+    )
+
+
 @decl.statement(dialect=dialect)
 class SiteBusMove(ir.Statement):
     y_mask: ir.SSAValue = info.argument(ilist.IListType[types.Bool, types.Literal(5)])

src/bloqade/lanes/dialects/move.py

@@ -10,10 +10,21 @@
 dialect = ir.Dialect(name="lanes.move")
 
 
+@statement(dialect=dialect)
+class Fill(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+
+    location_addresses: tuple[LocationAddress, ...] = info.attribute()
+
+
 @statement(dialect=dialect)
 class Initialize(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
+
     location_addresses: tuple[LocationAddress, ...] = info.attribute()
+    thetas: tuple[ir.SSAValue, ...] = info.argument(type=types.Float)
+    phis: tuple[ir.SSAValue, ...] = info.argument(type=types.Float)
+    lams: tuple[ir.SSAValue, ...] = info.argument(type=types.Float)
 
 
 @statement(dialect=dialect)

src/bloqade/lanes/dialects/place.py

@@ -30,6 +30,15 @@ def __init__(
         )
 
 
+@statement(dialect=dialect)
+class Initialize(QuantumStmt):
+    qubits: tuple[int, ...] = info.attribute()
+
+    theta: ir.SSAValue = info.argument(type=types.Float)
+    phi: ir.SSAValue = info.argument(type=types.Float)
+    lam: ir.SSAValue = info.argument(type=types.Float)
+
+
 @statement(dialect=dialect)
 class CZ(QuantumStmt):
     targets: tuple[int, ...] = info.attribute()
@@ -215,7 +224,7 @@ def end_measure(
         return (new_state,) + (EmptyLattice.bottom(),) * len(stmt.qubits)
 
 
-@dialect.register(key="circuit.layout")
+@dialect.register(key="place.layout")
 class InitialLayoutMethods(interp.MethodTable):
 
     @interp.impl(CZ)
@@ -229,6 +238,21 @@ def cz(
 
         return ()
 
+    @interp.impl(Initialize)
+    def initialize(
+        self,
+        _interp: LayoutAnalysis,
+        frame: ForwardFrame[EmptyLattice],
+        stmt: Initialize,
+    ):
+        for qubit_id in stmt.qubits:
+            global_qubit_id = _interp.global_address_stack[qubit_id]
+            _interp.thetas[global_qubit_id] = stmt.theta
+            _interp.phis[global_qubit_id] = stmt.phi
+            _interp.lams[global_qubit_id] = stmt.lam
+
+        return ()
+
     @interp.impl(StaticPlacement)
     def static_circuit(
         self,

src/bloqade/lanes/rewrite/place2move.py

@@ -292,6 +292,30 @@ def rewrite_Statement(self, node: ir.Statement):
 
 @dataclass
 class InsertInitialize(RewriteRule):
+    init_locations: tuple[LocationAddress, ...]
+    thetas: tuple[ir.SSAValue, ...]
+    phis: tuple[ir.SSAValue, ...]
+    lams: tuple[ir.SSAValue, ...]
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not (len(self.init_locations) > 0 and isinstance(node, func.Function)):
+            return RewriteResult()
+
+        first_stmt = node.body.blocks[0].first_stmt
+
+        if first_stmt is None or isinstance(first_stmt, move.Initialize):
+            return RewriteResult()
+        move.Initialize(
+            location_addresses=self.init_locations,
+            thetas=self.thetas,
+            phis=self.phis,
+            lams=self.lams,
+        ).insert_before(first_stmt)
+        return RewriteResult(has_done_something=True)
+
+
+@dataclass
+class InsertFill(RewriteRule):
     initial_layout: tuple[LocationAddress, ...]
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
@@ -300,11 +324,11 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         first_stmt = node.body.blocks[0].first_stmt
 
-        if first_stmt is None or isinstance(first_stmt, move.Initialize):
+        if first_stmt is None or isinstance(first_stmt, move.Fill):
             return RewriteResult()
-        move.Initialize(location_addresses=self.initial_layout).insert_before(
-            first_stmt
-        )
+
+        move.Fill(location_addresses=self.initial_layout).insert_before(first_stmt)
+
         return RewriteResult(has_done_something=True)
 
 

src/bloqade/lanes/upstream.py

@@ -52,9 +52,10 @@ def emit(self, mt: Method, no_raise: bool = True):
 
         if no_raise:
             address_frame, _ = address.AddressAnalysis(out.dialects).run_no_raise(out)
-            initial_layout = layout.LayoutAnalysis(
+            initial_layout, init_locations, thetas, phis, lams = layout.LayoutAnalysis(
                 out.dialects, self.layout_heristic, address_frame.entries
             ).get_layout_no_raise(out)
+
             placement_frame, _ = placement.PlacementAnalysis(
                 out.dialects,
                 initial_layout,
@@ -63,7 +64,7 @@ def emit(self, mt: Method, no_raise: bool = True):
             ).run_no_raise(out)
         else:
             address_frame, _ = address.AddressAnalysis(out.dialects).run(out)
-            initial_layout = layout.LayoutAnalysis(
+            initial_layout, init_locations, thetas, phis, lams = layout.LayoutAnalysis(
                 out.dialects, self.layout_heristic, address_frame.entries
             ).get_layout(out)
             placement_frame, _ = placement.PlacementAnalysis(
@@ -73,15 +74,14 @@ def emit(self, mt: Method, no_raise: bool = True):
                 self.placement_strategy,
             ).run(out)
 
-        placement_analysis = placement_frame.entries
-        args = (self.move_scheduler, placement_analysis)
         rule = rewrite.Chain(
-            place2move.InsertInitialize(initial_layout),
-            place2move.InsertMoves(*args),
-            place2move.RewriteCZ(*args),
-            place2move.RewriteR(*args),
-            place2move.RewriteRz(*args),
-            place2move.InsertMeasure(*args),
+            place2move.InsertFill(initial_layout),
+            place2move.InsertInitialize(init_locations, thetas, phis, lams),
+            place2move.InsertMoves(self.move_scheduler, placement_frame.entries),
+            place2move.RewriteCZ(self.move_scheduler, placement_frame.entries),
+            place2move.RewriteR(self.move_scheduler, placement_frame.entries),
+            place2move.RewriteRz(self.move_scheduler, placement_frame.entries),
+            place2move.InsertMeasure(self.move_scheduler, placement_frame.entries),
         )
         rewrite.Walk(rule).rewrite(out.code)