Adding gate metadata to spec + special field for non-static traps (#66)

Author: weinbe58

demo/log_depth_ghz.py

@@ -149,7 +149,12 @@ def run_ghz():
             static_traps={
                 "mem": mem,
             },
+            special_grid={
+                "gate": gate,
+            },
             fillable=set(["mem"]),
+            has_cz=set(["gate"]),
+            has_local=set(["mem"]),
         )
     )
 

src/bloqade/shuttle/arch.py

@@ -12,10 +12,27 @@ class Layout:
     """Abstract base class for layout."""
 
     fillable: set[str]
-    """ The set of trap names that are fillable by the sorter. """
+    """The set of trap names that are fillable by the sorter."""
+
+    has_cz: set[str]
+    """The set of trap names that can have a CZ gates applied."""
+
+    has_local: set[str]
+    """The set of trap names that can have local single qubit gates applied."""
+
+    special_grid: dict[str, Grid] = field(default_factory=dict, kw_only=True)
+    """Set of special grid values that are not static traps, but can be used for specific purposes."""
 
     def __hash__(self):
-        return hash((frozenset(self.static_traps.items()), frozenset(self.fillable)))
+        return hash(
+            (
+                frozenset(self.static_traps.items()),
+                frozenset(self.fillable),
+                frozenset(self.has_cz),
+                frozenset(self.has_local),
+                frozenset(self.special_grid.items()),
+            )
+        )
 
     def __eq__(self, other):
         if not isinstance(other, Layout):
@@ -81,7 +98,12 @@ def _default_layout():
         range(16),
     ).scale(10.0, 10.0)
 
-    return Layout(static_traps={"traps": zone}, fillable={"traps"})
+    return Layout(
+        static_traps={"traps": zone},
+        fillable={"traps"},
+        has_cz={"traps"},
+        has_local={"traps"},
+    )
 
 
 @dataclass(frozen=True)

src/bloqade/shuttle/stdlib/layouts/single_col_zone.py

@@ -27,8 +27,10 @@ def get_spec(num_x: int, num_y: int, spacing: float = 10.0) -> spec.ArchSpec:
 
     return spec.ArchSpec(
         layout=spec.Layout(
-            static_traps={"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
-            fillable=set(["traps"]),
+            {"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
+            set(["traps"]),
+            set(["traps"]),
+            set(["traps"]),
         )
     )
 

src/bloqade/shuttle/stdlib/layouts/two_col_zone.py

@@ -31,8 +31,10 @@ def get_spec(
 
     return spec.ArchSpec(
         layout=spec.Layout(
-            static_traps={"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
-            fillable=set(["traps"]),
+            {"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
+            set(["traps"]),
+            set(["traps"]),
+            set(["traps"]),
         )
     )
 

src/bloqade/shuttle/stdlib/spec.py

@@ -33,5 +33,7 @@ def single_zone_spec(num_x: int, num_y: int, spacing: float = 10.0) -> spec.Arch
         layout=spec.Layout(
             static_traps={"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
             fillable=set(["traps"]),
+            has_cz=set(["traps"]),
+            has_local=set(["traps"]),
         )
     )

test/analysis/test_zone.py

@@ -19,6 +19,8 @@ def default_spec():
         spec.Layout(
             static_traps={"test": grid.Grid.from_positions([0, 1, 2], [0, 1, 2])},
             fillable=set(["test"]),
+            has_cz=set(["test"]),
+            has_local=set(["test"]),
         )
     )
 

test/passes/test_inject_spec.py

@@ -5,11 +5,20 @@
 from bloqade.shuttle.prelude import move
 
 
+def get_spec(slm_grid: grid.Grid) -> spec.ArchSpec:
+    return spec.ArchSpec(
+        layout=spec.Layout(
+            {"slm": slm_grid},
+            fillable=set(["slm"]),
+            has_cz=set(["slm"]),
+            has_local=set(["slm"]),
+        )
+    )
+
+
 def test_inject_spec():
     slm_grid = grid.Grid.from_positions([1, 2, 3], [4, 5, 6])
-    test_spec = spec.ArchSpec(
-        layout=spec.Layout({"slm": slm_grid}, fillable=set(["slm"]))
-    )
+    test_spec = get_spec(slm_grid)
 
     @move(arch_spec=test_spec)
     def test():
@@ -21,11 +30,8 @@ def test():
 
 
 def test_inject_spac_callgraph():
-
     slm_grid = grid.Grid.from_positions([1, 2, 3], [4, 5, 6])
-    test_spec = spec.ArchSpec(
-        layout=spec.Layout({"slm": slm_grid}, fillable=set(["slm"]))
-    )
+    test_spec = get_spec(slm_grid)
 
     @move
     def subroutine(depth: int):

test/stdlib/test_spec.py

@@ -18,6 +18,8 @@ def test_single_zone_spec():
                 )
             },
             fillable=set(["traps"]),
+            has_cz=set(["traps"]),
+            has_local=set(["traps"]),
         )
     )
 

test/unit/test_arch.py

@@ -5,12 +5,26 @@
 
 def test_layout():
 
-    layout = Layout({"test": Grid.from_positions(range(16), range(16))}, {"test"})
+    layout = Layout(
+        {"test": Grid.from_positions(range(16), range(16))},
+        {"test"},
+        {"test"},
+        {"test"},
+    )
 
     assert hash(layout) == hash(
-        (frozenset(layout.static_traps.items()), frozenset(layout.fillable))
+        (
+            frozenset(layout.static_traps.items()),
+            frozenset(layout.fillable),
+            frozenset(layout.has_cz),
+            frozenset(layout.has_local),
+            frozenset(layout.special_grid.items()),
+        )
     )
     assert layout == Layout(
-        {"test": Grid.from_positions(range(16), range(16))}, {"test"}
+        {"test": Grid.from_positions(range(16), range(16))},
+        {"test"},
+        {"test"},
+        {"test"},
     )
     assert layout != 1