fixing bug related to permutation of position order in bloqade-geometry (#53)

Author: weinbe58

src/bloqade/lanes/analysis/atom/impl.py

@@ -170,9 +170,7 @@ def noop_impl(
     ):
         frame.set_state_for_stmt(stmt)
 
-    @interp.impl(move.Initialize)
-    def initialize_impl(
-        self, interp_: AtomInterpreter, frame: AtomFrame, stmt: move.Initialize
-    ):
+    @interp.impl(move.Fill)
+    def fill_impl(self, interp_: AtomInterpreter, frame: AtomFrame, stmt: move.Fill):
         frame.current_state = AtomState(stmt.location_addresses)
         frame.set_state_for_stmt(stmt)

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

@@ -6,8 +6,8 @@
 from bloqade.lanes.layout.word import Word
 
 
-def site_buses(word_size_y: int):
-    site_addresses = np.arange(word_size_y * 2).reshape((word_size_y, 2))
+def site_buses(site_addresses: np.ndarray):
+    word_size_y = site_addresses.shape[0]
 
     site_buses: list[Bus] = []
     for shift in range(word_size_y):
@@ -26,7 +26,6 @@ def site_buses(word_size_y: int):
                 src=as_flat_tuple_int(site_addresses[shift:, 0]),
             )
         )
-
     return tuple(site_buses)
 
 
@@ -62,18 +61,24 @@ def generate_arch(hypercube_dims: int = 4, word_size_y: int = 5) -> ArchSpec:
         Word(tuple(grid.shift(10.0 * ix, 0.0).positions)) for ix in range(num_word_x)
     )
 
-    word_ids = np.arange(word_size_x * word_size_y).reshape(word_size_y, word_size_x)
+    site_ids = (
+        np.arange(word_size_x * word_size_y)
+        .reshape(word_size_x, word_size_y)
+        .transpose()
+    )
     word_buses = hypercube_busses(hypercube_dims)
     site_bus_compatibility = tuple(
         frozenset(range(num_word_x)) for _ in range(num_word_x)
     )
+
     gate_zone = tuple(range(len(words)))
+
     return ArchSpec(
         words,
         (gate_zone,),
         frozenset(range(num_word_x)),
-        frozenset(as_flat_tuple_int(word_ids[:, 1])),
-        site_buses(word_size_y),
+        frozenset(as_flat_tuple_int(site_ids[:, 1])),
+        site_buses(site_ids),
         word_buses,
         site_bus_compatibility,
     )

src/bloqade/lanes/heuristics/fixed.py

@@ -46,9 +46,9 @@ def validate_initial_layout(
                 raise ValueError(
                     "Initial layout contains invalid word id for logical arch"
                 )
-            if addr.site_id % 2 != 0:
+            if addr.site_id >= 5:
                 raise ValueError(
-                    "Initial layout should only contain even site ids for fixed home location strategy"
+                    "Initial layout should only site ids < 5 for logical arch"
                 )
 
     def _word_balance(
@@ -125,7 +125,7 @@ def cz_placements(
             mover, dst_addr = self._pick_mover_and_location(state, start_word_id, c, t)
             new_positions[mover] = LocationAddress(
                 word_id=dst_addr.word_id,
-                site_id=dst_addr.site_id + 1,
+                site_id=dst_addr.site_id + 5,
             )
 
         return self._update_positions(state, new_positions)
@@ -158,6 +158,9 @@ def assert_valid_word_bus_move(
         assert (
             src_site in self.arch_spec.has_word_buses
         ), f"Invalid source site {src_site} for word bus move"
+        assert src_word < len(
+            self.arch_spec.words
+        ), f"Invalid source word {src_word} for site bus move {bus_id}"
 
         return WordLaneAddress(
             direction,
@@ -182,6 +185,9 @@ def assert_valid_site_bus_move(
         assert (
             src_word in self.arch_spec.has_site_buses
         ), f"Invalid source word {src_word} for site bus move {bus_id}"
+        assert src_word < len(
+            self.arch_spec.words
+        ), f"Invalid source word {src_word} for site bus move {bus_id}"
 
         return SiteLaneAddress(
             direction,
@@ -200,7 +206,8 @@ def site_moves(
 
         bus_moves = {}
         for before, end in diffs:
-            bus_id = end.site_id // 2 - before.site_id // 2
+            bus_id = (end.site_id % 5) - (before.site_id % 5)
+
             if bus_id < 0:
                 bus_id += len(self.arch_spec.site_buses)
 
@@ -265,6 +272,9 @@ def compute_moves(
         moves.extend(self.site_moves(groups.get((0, 0), []), 0))
         moves.extend(self.site_moves(groups.get((1, 1), []), 1))
 
+        for move_group in moves:
+            print(move_group)
+
         return moves
 
 
@@ -322,9 +332,8 @@ def compute_layout(
         available_addresses = set(
             [
                 LocationAddress(word_id, site_id)
-                for word_id, word in enumerate(self.arch_spec.words)
-                for site_id in range(len(word.sites))
-                if site_id % 2 == 0
+                for word_id in range(len(self.arch_spec.words))
+                for site_id in range(5)
             ]
         )
 

src/bloqade/lanes/visualize.py

@@ -133,6 +133,8 @@ def get_drawer(mt: ir.Method, arch_spec: ArchSpec, ax: Axes, atom_marker: str =
             steps.append((stmt, curr_state))
 
     def draw(step_index: int):
+        if len(steps) == 0:
+            return
         stmt, curr_state = steps[step_index]
         show_slm(ax, stmt, arch_spec, atom_marker)
 

test/heuristics/test_fixed.py

@@ -30,19 +30,19 @@ def cz_placement_cases():
         occupied=frozenset(),
         layout=(
             LocationAddress(0, 0),
-            LocationAddress(0, 2),
+            LocationAddress(0, 1),
             LocationAddress(1, 0),
-            LocationAddress(1, 2),
+            LocationAddress(1, 1),
         ),
         move_count=(0, 0, 0, 0),
     )
     state_after = ConcreteState(
         occupied=frozenset(),
         layout=(
-            LocationAddress(1, 1),
-            LocationAddress(1, 3),
+            LocationAddress(1, 5),
+            LocationAddress(1, 6),
             LocationAddress(1, 0),
-            LocationAddress(1, 2),
+            LocationAddress(1, 1),
         ),
         move_count=(1, 1, 0, 0),
     )
@@ -58,19 +58,19 @@ def cz_placement_cases():
         occupied=frozenset(),
         layout=(
             LocationAddress(0, 0),
-            LocationAddress(0, 2),
+            LocationAddress(0, 1),
             LocationAddress(1, 0),
-            LocationAddress(1, 2),
+            LocationAddress(1, 1),
         ),
         move_count=(1, 1, 0, 0),
     )
     state_after = ConcreteState(
         occupied=frozenset(),
         layout=(
             LocationAddress(0, 0),
-            LocationAddress(0, 2),
             LocationAddress(0, 1),
-            LocationAddress(0, 3),
+            LocationAddress(0, 5),
+            LocationAddress(0, 6),
         ),
         move_count=(1, 1, 1, 1),
     )
@@ -85,19 +85,19 @@ def cz_placement_cases():
         occupied=frozenset(),
         layout=(
             LocationAddress(0, 0),
+            LocationAddress(0, 1),
             LocationAddress(0, 2),
-            LocationAddress(0, 4),
-            LocationAddress(0, 6),
+            LocationAddress(0, 3),
         ),
         move_count=(1, 1, 0, 0),
     )
     state_after = ConcreteState(
         occupied=frozenset(),
         layout=(
             LocationAddress(0, 0),
-            LocationAddress(0, 2),
             LocationAddress(0, 1),
-            LocationAddress(0, 3),
+            LocationAddress(0, 5),
+            LocationAddress(0, 6),
         ),
         move_count=(1, 1, 1, 1),
     )
@@ -112,19 +112,19 @@ def cz_placement_cases():
         occupied=frozenset(),
         layout=(
             LocationAddress(0, 0),
+            LocationAddress(0, 1),
             LocationAddress(0, 2),
-            LocationAddress(0, 4),
-            LocationAddress(0, 6),
+            LocationAddress(0, 3),
         ),
         move_count=(0, 0, 1, 1),
     )
     state_after = ConcreteState(
         occupied=frozenset(),
         layout=(
-            LocationAddress(0, 5),
             LocationAddress(0, 7),
-            LocationAddress(0, 4),
-            LocationAddress(0, 6),
+            LocationAddress(0, 8),
+            LocationAddress(0, 2),
+            LocationAddress(0, 3),
         ),
         move_count=(1, 1, 1, 1),
     )
@@ -185,7 +185,7 @@ def test_fixed_invalid_initial_layout_1():
         LocationAddress(0, 0),
         LocationAddress(0, 1),
         LocationAddress(0, 2),
-        LocationAddress(0, 3),
+        LocationAddress(0, 5),
     )
     with pytest.raises(ValueError):
         placement_strategy.validate_initial_layout(layout)
@@ -215,15 +215,15 @@ def test_initial_layout():
     print(layout)
     assert layout == (
         LocationAddress(word_id=0, site_id=0),
+        LocationAddress(word_id=0, site_id=1),
         LocationAddress(word_id=0, site_id=2),
+        LocationAddress(word_id=0, site_id=3),
         LocationAddress(word_id=0, site_id=4),
-        LocationAddress(word_id=0, site_id=6),
-        LocationAddress(word_id=0, site_id=8),
         LocationAddress(word_id=1, site_id=0),
+        LocationAddress(word_id=1, site_id=1),
         LocationAddress(word_id=1, site_id=2),
+        LocationAddress(word_id=1, site_id=3),
         LocationAddress(word_id=1, site_id=4),
-        LocationAddress(word_id=1, site_id=6),
-        LocationAddress(word_id=1, site_id=8),
     )
 
 
@@ -234,7 +234,7 @@ def test_move_scheduler_cz():
         tuple(
             LocationAddress(word_id, site_id)
             for word_id in range(2)
-            for site_id in range(0, 10, 2)
+            for site_id in range(5)
         ),
         tuple(0 for _ in range(10)),
     )
@@ -257,19 +257,19 @@ def test_move_scheduler_cz():
                 direction=Direction.FORWARD, word_id=0, site_id=0, bus_id=0
             ),
             SiteLaneAddress(
-                direction=Direction.FORWARD, word_id=0, site_id=2, bus_id=0
+                direction=Direction.FORWARD, word_id=0, site_id=1, bus_id=0
             ),
         ),
-        (SiteLaneAddress(direction=Direction.FORWARD, word_id=0, site_id=8, bus_id=7),),
+        (SiteLaneAddress(direction=Direction.FORWARD, word_id=0, site_id=4, bus_id=7),),
         (
             WordLaneAddress(
-                direction=Direction.FORWARD, word_id=0, site_id=1, bus_id=0
+                direction=Direction.FORWARD, word_id=0, site_id=5, bus_id=0
             ),
             WordLaneAddress(
-                direction=Direction.FORWARD, word_id=0, site_id=3, bus_id=0
+                direction=Direction.FORWARD, word_id=0, site_id=6, bus_id=0
             ),
             WordLaneAddress(
-                direction=Direction.FORWARD, word_id=0, site_id=5, bus_id=0
+                direction=Direction.FORWARD, word_id=0, site_id=7, bus_id=0
             ),
         ),
     ]