adding auto encoding

Author: weinbe58

src/bloqade/lanes/types/arch.py

@@ -1,8 +1,10 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Generic, Sequence
 
 import numpy as np
 
+from bloqade.lanes.types.encoding import EncodingType
+
 from .word import SiteType, Word
 
 
@@ -33,6 +35,10 @@ class ArchSpec(Generic[SiteType]):
     """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."""
+    encoding: EncodingType = field(init=False)
+
+    def __post_init__(self):
+        object.__setattr__(self, "encoding", EncodingType.infer(self))  # type: ignore
 
     def plot(
         self,

src/bloqade/lanes/types/encoding.py

@@ -17,6 +17,27 @@ class EncodingType(enum.Enum):
     BIT32 = 0
     BIT64 = 1
 
+    @staticmethod
+    def infer(spec) -> "EncodingType":
+        num_words = len(spec.words)
+        num_sites = len(spec.words[0].sites)
+        num_site_buses = len(spec.site_buses)
+        num_word_buses = len(spec.word_buses)
+
+        max_id = max(
+            num_words - 1,
+            num_sites - 1,
+            num_site_buses - 1,
+            num_word_buses - 1,
+        )
+
+        if max_id < 256:
+            return EncodingType.BIT32
+        elif max_id < 65536:
+            return EncodingType.BIT64
+        else:
+            raise ValueError("Architecture too large to encode with 64-bit addresses")
+
 
 @dataclass(frozen=True)
 class MoveType(abc.ABC):

src/bloqade/lanes/types/path.py

@@ -5,7 +5,7 @@
 import rustworkx as nx
 
 from .arch import ArchSpec
-from .encoding import Direction, EncodingType, InterMove, IntraMove, MoveType
+from .encoding import Direction, InterMove, IntraMove, MoveType
 
 
 @dataclass(frozen=True)
@@ -56,7 +56,7 @@ def __post_init__(self):
                         lane_addr,
                     )
 
-    def extract_lanes_from_path(self, path: list[int], encoding: EncodingType):
+    def extract_lanes_from_path(self, path: list[int]):
         """Given a path as a list of node indices, extract the lane addresses."""
         if len(path) < 2:
             raise ValueError("Path must have at least two nodes to extract lanes.")
@@ -67,15 +67,14 @@ def extract_lanes_from_path(self, path: list[int], encoding: EncodingType):
 
             lane: MoveType = self.site_graph.get_edge_data(src, dst)
 
-            lanes.append(lane.get_address(encoding))
+            lanes.append(lane.get_address(self.spec.encoding))
         return lanes
 
     def find_path(
         self,
         start: tuple[int, int],
         end: tuple[int, int],
         occupied: frozenset[tuple[int, int]] = frozenset(),
-        encoding: EncodingType = EncodingType.BIT64,
         path_heuristic: Callable[[list[tuple[int, int]]], float] = lambda _: 0.0,
     ):
         """Find a path from start to end avoiding occupied sites.
@@ -114,7 +113,7 @@ def find_path(
             path_nodes,
             key=lambda p: path_heuristic([self.physical_addresses[n] for n in p]),
         )
-        lanes = self.extract_lanes_from_path(path, encoding)
+        lanes = self.extract_lanes_from_path(path)
         return lanes, occupied.union(
             frozenset(self.physical_addresses[n] for n in path)
         )