updating variables names to match naming convention + adding site_bus_compatibility to arch

Author: weinbe58

src/bloqade/lanes/gemini.py

@@ -1,65 +1,67 @@
 import numpy as np
 
-from bloqade.lanes.types.arch import ArchSpec, Lane
-from bloqade.lanes.types.block import Block
+from bloqade.lanes.types.arch import ArchSpec, Bus
 from bloqade.lanes.types.grid import Grid
 from bloqade.lanes.types.numpy_compat import as_flat_tuple_int
+from bloqade.lanes.types.word import Word
 
 
-def intra_moves(block_size_y: int):
-    intra_ids = np.arange(block_size_y * 2).reshape((block_size_y, 2))
+def site_buses(word_size_y: int):
+    site_addresses = np.arange(word_size_y * 2).reshape((word_size_y, 2))
 
-    intra_lanes: list[Lane] = []
-    for shift in range(block_size_y):
-        intra_lanes.append(
-            Lane(
-                src=as_flat_tuple_int(intra_ids[: block_size_y - shift, 0]),
-                dst=as_flat_tuple_int(intra_ids[shift:, 1]),
+    site_buses: list[Bus] = []
+    for shift in range(word_size_y):
+        site_buses.append(
+            Bus(
+                src=as_flat_tuple_int(site_addresses[: word_size_y - shift, 0]),
+                dst=as_flat_tuple_int(site_addresses[shift:, 1]),
             )
         )
 
-    return tuple(intra_lanes)
+    return tuple(site_buses)
 
 
 def generate_arch():
 
-    block_size_y = 5
-    block_size_x = 2
-    num_block_x = 16
+    word_size_y = 5
+    word_size_x = 2
+    num_word_x = 16
 
     x_positions = (0.0, 2.0)
-    y_positions = tuple(10.0 * i for i in range(block_size_y))
+    y_positions = tuple(10.0 * i for i in range(word_size_y))
 
     grid = Grid(x_positions, y_positions)
 
-    blocks = tuple(
-        Block(grid.shift(10.0 * ix, 0.0).positions) for ix in range(num_block_x)
+    words = tuple(
+        Word(grid.shift(10.0 * ix, 0.0).positions) for ix in range(num_word_x)
     )
 
-    inter_lanes: list[Lane] = []
+    word_buses: list[Bus] = []
     for shift in range(4):
         m = 1 << shift
 
         srcs = []
         dsts = []
-        for src in range(num_block_x):
+        for src in range(num_word_x):
             if src & m != 0:
                 continue
 
             dst = src | m
             srcs.append(src)
             dsts.append(dst)
 
-        inter_lanes.append(Lane(tuple(srcs), tuple(dsts)))
+        word_buses.append(Bus(tuple(srcs), tuple(dsts)))
 
-    block_ids = np.arange(block_size_x * block_size_y).reshape(
-        block_size_y, block_size_x
-    )
+    word_ids = np.arange(word_size_x * word_size_y).reshape(word_size_y, word_size_x)
 
+    site_bus_compatibility = tuple(
+        frozenset(j for j in range(num_word_x) if j != i) for i in range(num_word_x)
+    )
     return ArchSpec(
-        blocks=blocks,
-        has_intra_lanes=frozenset(range(num_block_x)),
-        has_inter_lanes=frozenset(as_flat_tuple_int(block_ids[:, 1])),
-        intra_lanes=intra_moves(block_size_y),
-        inter_lanes=tuple(inter_lanes),
+        words,
+        frozenset(range(num_word_x)),
+        frozenset(as_flat_tuple_int(word_ids[:, 1])),
+        site_buses(word_size_y),
+        tuple(word_buses),
+        site_bus_compatibility,
     )

src/bloqade/lanes/types/arch.py

@@ -1,18 +1,16 @@
 from dataclasses import dataclass
-from itertools import chain
 from typing import Generic, Sequence
 
 import numpy as np
-import rustworkx as nx
 
-from .block import Block, SiteType
+from .word import SiteType, Word
 
 
 @dataclass(frozen=True)
-class Lane:
+class Bus:
     """A group of inter-lanes that can be executed in parallel.
 
-    For inter-lanes, src and dst are the block indices involved in the inter-lane.
+    For inter-lanes, src and dst are the word indices involved in the inter-lane.
     For intra-lanes, src are the source site indices and dst are the destination site indices.
 
     """
@@ -23,55 +21,23 @@ class Lane:
 
 @dataclass(frozen=True)
 class ArchSpec(Generic[SiteType]):
-    blocks: tuple[Block[SiteType], ...]
-    """List of all blocks in the architecture."""
-    has_intra_lanes: frozenset[int]
-    """Set of blocks that have intra-lane moves."""
-    has_inter_lanes: frozenset[int]  # set of sites in block that have inter-lanes moves
-    """Set of sites (by index) that have inter-lane moves. These sites are the same across all blocks."""
-    intra_lanes: tuple[Lane, ...]
-    """List of all intra-lanes in the architecture by site index."""
-    inter_lanes: tuple[Lane, ...]
-    """List of all inter-lanes in the architecture by block index."""
-
-    def get_graphs(self):
-
-        block_intra_graph = nx.PyDiGraph()
-        intra_move_compat = nx.PyGraph()
-
-        all_moves = list(
-            chain.from_iterable(
-                (list(zip(lane.src, lane.dst)) for lane in self.intra_lanes)
-            )
-        )
-
-        move_map = {move: i for i, move in enumerate(all_moves)}
-
-        for lane in self.intra_lanes:
-            moves = list(zip(lane.src, lane.dst))
-            for src, dst in moves:
-                block_intra_graph.add_edge(src, dst, None)
-
-            for i, move_i in enumerate(moves):
-                for j, move_j in enumerate(moves[i + 1 :], i + 1):
-                    intra_move_compat.add_edge(move_map[move_i], move_map[move_j], None)
-
-        inter_block_graph = nx.PyDiGraph()
-        inter_move_compat = nx.PyGraph()
-
-        for lane in self.inter_lanes:
-            moves = list(zip(lane.src, lane.dst))
-            for src, dst in moves:
-                inter_block_graph.add_edge(src, dst, None)
-
-            for i, move_i in enumerate(moves):
-                for j, move_j in enumerate(moves[i + 1 :], i + 1):
-                    inter_move_compat.add_edge(i, j, None)
+    words: tuple[Word[SiteType], ...]
+    """tuple of all words in the architecture. words[i] gives the word at word address i."""
+    has_site_buses: frozenset[int]
+    """Set of words that have site-lane moves."""
+    has_word_buses: frozenset[int]  # set of sites in word that have inter-lanes moves
+    """Set of sites (by index) that have word-lane moves. These sites are the same across all words."""
+    site_buses: tuple[Bus, ...]
+    """List of all site buses in the architecture by site address."""
+    word_buses: tuple[Bus, ...]
+    """List of all word buses in the architecture by word address."""
+    site_bus_compatibility: tuple[frozenset[int], ...]
+    """Mapping from word id indicating which other word ids can execute site-buses in parallel."""
 
     def plot(
         self,
         ax=None,
-        show_blocks: Sequence[int] = (),
+        show_words: Sequence[int] = (),
         show_intra: Sequence[int] = (),
         show_inter: Sequence[int] = (),
         **scatter_kwargs,
@@ -82,9 +48,9 @@ def plot(
         if ax is None:
             ax = plt.gca()
 
-        for block_id in show_blocks:
-            block = self.blocks[block_id]
-            block.plot(ax)
+        for word_id in show_words:
+            word = self.words[word_id]
+            word.plot(ax)
 
         x_min, x_max = ax.get_xlim()
         y_min, y_max = ax.get_ylim()
@@ -93,14 +59,14 @@ def plot(
         bow_x = (x_max - x_min) * 0.025
 
         colors = {}
-        for block_id in show_blocks:
-            block = self.blocks[block_id]
+        for word_id in show_words:
+            word = self.words[word_id]
             for lane_id in show_intra:
-                lane = self.intra_lanes[lane_id]
+                lane = self.site_buses[lane_id]
 
                 for start, end in zip(lane.src, lane.dst):
-                    start = block[start]
-                    end = block[end]
+                    start = word[start]
+                    end = word[end]
 
                     for (x_start, y_start), (x_end, y_end) in zip(
                         start.positions(), end.positions()
@@ -127,14 +93,14 @@ def plot(
                             colors[lane] = ln.get_color()
 
         for lane in show_inter:
-            lane = self.inter_lanes[lane]
-            for start_block_id, end_block_id in zip(lane.src, lane.dst):
-                start_block = self.blocks[start_block_id]
-                end_block = self.blocks[end_block_id]
-
-                for site in self.has_inter_lanes:
-                    start = start_block[site]
-                    end = end_block[site]
+            lane = self.word_buses[lane]
+            for start_word_id, end_word_id in zip(lane.src, lane.dst):
+                start_word = self.words[start_word_id]
+                end_word = self.words[end_word_id]
+
+                for site in self.has_word_buses:
+                    start = start_word[site]
+                    end = end_word[site]
                     for (x_start, y_start), (x_end, y_end) in zip(
                         start.positions(), end.positions()
                     ):
@@ -164,7 +130,7 @@ def plot(
     def show(
         self,
         ax=None,
-        show_blocks: Sequence[int] = (),
+        show_words: Sequence[int] = (),
         show_intra: Sequence[int] = (),
         show_inter: Sequence[int] = (),
         **scatter_kwargs,
@@ -173,7 +139,7 @@ def show(
 
         self.plot(
             ax,
-            show_blocks=show_blocks,
+            show_words=show_words,
             show_intra=show_intra,
             show_inter=show_inter,
             **scatter_kwargs,

src/bloqade/lanes/types/encoding.py

@@ -41,12 +41,12 @@ def reverse(self):
 
 @dataclass(frozen=True)
 class IntraMove(MoveType):
-    block_id: int
+    word_id: int
     site_id: int
     lane_id: int
 
     def src_site(self) -> tuple[int, int]:
-        return self.block_id, self.site_id
+        return self.word_id, self.site_id
 
     def get_address(self, encoding: EncodingType) -> int:
         if encoding == EncodingType.BIT32:
@@ -60,17 +60,17 @@ def get_address(self, encoding: EncodingType) -> int:
         else:
             raise ValueError("Unsupported encoding type")
 
-        block_id_enc = mask & self.block_id
+        word_id_enc = mask & self.word_id
         lane_id_enc = mask & self.lane_id
         site_id_enc = mask & self.site_id
 
-        assert block_id_enc == self.block_id, "Block ID too large to encode"
+        assert word_id_enc == self.word_id, "word ID too large to encode"
         assert lane_id_enc == self.lane_id, "Lane ID too large to encode"
         assert site_id_enc == self.site_id, "Site ID too large to encode"
 
         address = lane_id_enc
         address |= site_id_enc << shift
-        address |= block_id_enc << (2 * shift)
+        address |= word_id_enc << (2 * shift)
         address |= MoveTypeEnum.INTRA.value << (3 * shift + padding)
         address |= self.direction.value << (3 * shift + padding + 1)
         assert address.bit_length() <= (
@@ -81,12 +81,12 @@ def get_address(self, encoding: EncodingType) -> int:
 
 @dataclass(frozen=True)
 class InterMove(MoveType):
-    start_block_id: int
-    end_block_id: int
+    start_word_id: int
+    end_word_id: int
     lane_id: int
 
     def src_site(self) -> tuple[int, int]:
-        return self.start_block_id, self.lane_id
+        return self.start_word_id, self.lane_id
 
     def get_address(self, encoding: EncodingType) -> int:
         if encoding == EncodingType.BIT32:
@@ -101,18 +101,16 @@ def get_address(self, encoding: EncodingType) -> int:
             raise ValueError("Unsupported encoding type")
 
         lane_id_enc = mask & self.lane_id
-        start_block_id = mask & self.start_block_id
-        end_block_id = mask & self.end_block_id
+        start_word_id = mask & self.start_word_id
+        end_word_id = mask & self.end_word_id
 
         assert lane_id_enc == self.lane_id, "Lane ID too large to encode"
-        assert (
-            start_block_id == self.start_block_id
-        ), "Start Block ID too large to encode"
-        assert end_block_id == self.end_block_id, "End Block ID too large to encode"
+        assert start_word_id == self.start_word_id, "Start word ID too large to encode"
+        assert end_word_id == self.end_word_id, "End word ID too large to encode"
 
         address = lane_id_enc
-        address |= end_block_id << shift
-        address |= start_block_id << (2 * shift)
+        address |= end_word_id << shift
+        address |= start_word_id << (2 * shift)
         address |= MoveTypeEnum.INTER.value << (3 * shift + padding)
         address |= self.direction.value << (3 * shift + padding + 1)
         assert address.bit_length() <= (