implement vacant grid (#31)

* implement vacant grid

* adding positions API

* moving position of property

* adding wrappers for filled

* refactor is_equal

* adding backup

* fixing names

* fixing comments from cduck

* Fixing bug in interpreter and adding test for repeat

* adding tests + fixing runtime

Author: weinbe58

src/bloqade/shuttle/__init__.py

@@ -1,5 +1,6 @@
 from .dialects.action import _interface as action
 from .dialects.atom import _interface as atom
+from .dialects.filled import _interface as filled
 from .dialects.gate import _interface as gate
 from .dialects.init import _interface as init
 from .dialects.measure import _interface as measure
@@ -15,4 +16,5 @@
     "measure",
     "schedule",
     "spec",
+    "filled",
 ]

src/bloqade/shuttle/dialects/filled/__init__.py

@@ -0,0 +1,9 @@
+from ._dialect import dialect as dialect
+from ._interface import (
+    fill as fill,
+    get_parent as get_parent,
+    vacate as vacate,
+)
+from .concrete import FilledGridMethods as FilledGridMethods
+from .stmts import Fill as Fill, GetParent as GetParent, Vacate as Vacate
+from .types import FilledGrid as FilledGrid, FilledGridType as FilledGridType

src/bloqade/shuttle/dialects/filled/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("shuttle.filled")

src/bloqade/shuttle/dialects/filled/_interface.py

@@ -0,0 +1,51 @@
+from typing import Any, TypeVar
+
+from bloqade.geometry.dialects import grid
+from kirin.dialects import ilist
+from kirin.lowering import wraps as _wraps
+
+from .stmts import Fill, GetParent, Vacate
+from .types import FilledGrid
+
+Nx = TypeVar("Nx")
+Ny = TypeVar("Ny")
+
+
+@_wraps(Vacate)
+def vacate(
+    zone: grid.Grid[Nx, Ny],
+    vacancies: ilist.IList[tuple[int, int], Any],
+) -> FilledGrid[Nx, Ny]: ...
+
+
+@_wraps(Fill)
+def fill(
+    zone: grid.Grid[Nx, Ny],
+    filled: ilist.IList[tuple[int, int], Any],
+) -> FilledGrid[Nx, Ny]: ...
+
+
+@_wraps(GetParent)
+def get_parent(filled_grid: FilledGrid[Nx, Ny]) -> grid.Grid[Nx, Ny]: ...
+
+
+@_wraps(grid.Shift)
+def shift(
+    grid: FilledGrid[Nx, Ny], x_shift: float, y_shift: float
+) -> FilledGrid[Nx, Ny]: ...
+
+
+@_wraps(grid.Scale)
+def scale(
+    grid: FilledGrid[Nx, Ny], x_scale: float, y_scale: float
+) -> FilledGrid[Nx, Ny]: ...
+
+
+@_wraps(grid.Repeat)
+def repeat(
+    grid: FilledGrid[Any, Any],
+    x_times: int,
+    y_times: int,
+    y_spacing: float,
+    x_spacing: float,
+) -> FilledGrid[Any, Any]: ...

src/bloqade/shuttle/dialects/filled/concrete.py

@@ -0,0 +1,35 @@
+from typing import Any
+
+from bloqade.geometry.dialects.grid.types import Grid
+from kirin.dialects import ilist
+from kirin.interp import (
+    Frame,
+    Interpreter,
+    MethodTable,
+    impl,
+)
+
+from . import stmts
+from ._dialect import dialect
+from .types import FilledGrid
+
+
+@dialect.register
+class FilledGridMethods(MethodTable):
+
+    @impl(stmts.Vacate)
+    def vacate(self, interp: Interpreter, frame: Frame, stmt: stmts.Vacate):
+        zone = frame.get_casted(stmt.zone, Grid)
+        vacancies = frame.get_casted(stmt.vacancies, ilist.IList[tuple[int, int], Any])
+        return (FilledGrid.vacate(zone, vacancies),)
+
+    @impl(stmts.Fill)
+    def fill(self, interp: Interpreter, frame: Frame, stmt: stmts.Fill):
+        zone = frame.get_casted(stmt.zone, Grid)
+        filled = frame.get_casted(stmt.filled, ilist.IList[tuple[int, int], Any])
+        return (FilledGrid.fill(zone, filled),)
+
+    @impl(stmts.GetParent)
+    def get_parent(self, interp: Interpreter, frame: Frame, stmt: stmts.GetParent):
+        filled_grid = frame.get_casted(stmt.filled_grid, FilledGrid)
+        return (filled_grid.parent,)

src/bloqade/shuttle/dialects/filled/stmts.py

@@ -0,0 +1,42 @@
+from bloqade.geometry.dialects import grid
+from kirin import decl, ir, lowering, types
+from kirin.decl import info
+from kirin.dialects import ilist
+
+from ._dialect import dialect
+from .types import FilledGridType
+
+NumVacant = types.TypeVar("NumVacant")
+Nx = types.TypeVar("Nx")
+Ny = types.TypeVar("Ny")
+
+
+@decl.statement(dialect=dialect)
+class Vacate(ir.Statement):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    zone: ir.SSAValue = info.argument(grid.GridType[Nx, Ny])
+    vacancies: ir.SSAValue = info.argument(
+        ilist.IListType[types.Tuple[types.Int, types.Int], NumVacant]
+    )
+    result: ir.ResultValue = info.result(FilledGridType[Nx, Ny])
+
+
+@decl.statement(dialect=dialect)
+class Fill(ir.Statement):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    zone: ir.SSAValue = info.argument(grid.GridType[Nx, Ny])
+    filled: ir.SSAValue = info.argument(
+        ilist.IListType[types.Tuple[types.Int, types.Int], NumVacant]
+    )
+    result: ir.ResultValue = info.result(FilledGridType[Nx, Ny])
+
+
+@decl.statement(dialect=dialect)
+class GetParent(ir.Statement):
+    name = "get_parent"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    filled_grid: ir.SSAValue = info.argument(FilledGridType[Nx, Ny])
+    result: ir.ResultValue = info.result(grid.GridType[Nx, Ny])

src/bloqade/shuttle/dialects/filled/types.py

@@ -0,0 +1,123 @@
+from dataclasses import dataclass, field
+from functools import cached_property
+from itertools import product
+from typing import Any, Iterable, Sequence, TypeVar
+
+from bloqade.geometry.dialects import grid
+from kirin import types
+from kirin.dialects import ilist
+
+NumX = TypeVar("NumX")
+NumY = TypeVar("NumY")
+
+
+@dataclass(eq=False)
+class FilledGrid(grid.Grid[NumX, NumY]):
+    x_spacing: tuple[float, ...] = field(init=False)
+    y_spacing: tuple[float, ...] = field(init=False)
+    x_init: float | None = field(init=False)
+    y_init: float | None = field(init=False)
+
+    parent: grid.Grid[NumX, NumY]
+    vacancies: frozenset[tuple[int, int]]
+
+    def __post_init__(self):
+        self.x_spacing = self.parent.x_spacing
+        self.y_spacing = self.parent.y_spacing
+        self.x_init = self.parent.x_init
+        self.y_init = self.parent.y_init
+
+        self.type = types.Generic(
+            FilledGrid,
+            types.Literal(len(self.x_spacing) + 1),
+            types.Literal(len(self.y_spacing) + 1),
+        )
+
+    def __hash__(self):
+        return hash((FilledGrid, self.parent, self.vacancies))
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            isinstance(other, FilledGrid)
+            and self.parent == other.parent
+            and self.vacancies == other.vacancies
+        )
+
+    @cached_property
+    def positions(self) -> ilist.IList[tuple[float, float], Any]:
+        positions = tuple(
+            (x, y)
+            for (ix, x), (iy, y) in product(
+                enumerate(self.x_positions), enumerate(self.y_positions)
+            )
+            if (ix, iy) not in self.vacancies
+        )
+
+        return ilist.IList(positions)
+
+    @classmethod
+    def fill(
+        cls, grid_obj: grid.Grid[NumX, NumY], filled: Sequence[tuple[int, int]]
+    ) -> "FilledGrid[NumX, NumY]":
+        num_x, num_y = grid_obj.shape
+
+        if isinstance(grid_obj, FilledGrid):
+            vacancies = grid_obj.vacancies
+        else:
+            vacancies = frozenset(product(range(num_x), range(num_y)))
+
+        vacancies = vacancies - frozenset(filled)
+
+        return cls(parent=grid_obj, vacancies=vacancies)
+
+    @classmethod
+    def vacate(
+        cls, grid_obj: grid.Grid[NumX, NumY], vacancies: Iterable[tuple[int, int]]
+    ) -> "FilledGrid[NumX, NumY]":
+
+        if isinstance(grid_obj, FilledGrid):
+            input_vacancies = grid_obj.vacancies
+        else:
+            input_vacancies = frozenset()
+
+        input_vacancies = input_vacancies.union(vacancies)
+
+        return cls(parent=grid_obj, vacancies=input_vacancies)
+
+    def get_view(  # type: ignore
+        self, x_indices: ilist.IList[int, Any], y_indices: ilist.IList[int, Any]
+    ):
+        remapping_x = {ix: i for i, ix in enumerate(x_indices)}
+        remapping_y = {iy: i for i, iy in enumerate(y_indices)}
+        return FilledGrid(
+            parent=self.parent.get_view(x_indices, y_indices),
+            vacancies=frozenset(
+                (remapping_x[x], remapping_y[y])
+                for x, y in self.vacancies
+                if x in remapping_x and y in remapping_y
+            ),
+        )
+
+    def shift(self, x_shift: float, y_shift: float):
+        return FilledGrid(
+            parent=self.parent.shift(x_shift, y_shift),
+            vacancies=self.vacancies,
+        )
+
+    def scale(self, x_scale: float, y_scale: float):
+        return FilledGrid(
+            parent=self.parent.scale(x_scale, y_scale),
+            vacancies=self.vacancies,
+        )
+
+    def repeat(self, x_times: int, y_times: int, x_gap: float, y_gap: float):
+        new_parent = self.parent.repeat(x_times, y_times, x_gap, y_gap)
+        x_dim, y_dim = self.shape
+        vacancies = frozenset(
+            (x + x_dim * i, y + y_dim * j)
+            for i, j, (x, y) in product(range(x_times), range(y_times), self.vacancies)
+        )
+        return FilledGrid.vacate(new_parent, vacancies)
+
+
+FilledGridType = types.Generic(FilledGrid, types.TypeVar("NumX"), types.TypeVar("NumY"))

src/bloqade/shuttle/prelude.py

@@ -11,6 +11,7 @@
 from bloqade.shuttle.dialects import (
     action,
     atom,
+    filled,
     gate,
     init,
     measure,
@@ -23,7 +24,7 @@
 from bloqade.shuttle.rewrite.desugar import DesugarTurnOffRewrite, DesugarTurnOnRewrite
 
 
-@ir.dialect_group(structural.union([spec, grid, atom, gate, op, qubit]))
+@ir.dialect_group(structural.union([spec, grid, filled, atom, gate, op, qubit]))
 def kernel(self):
     def run_pass(
         mt: ir.Method,
@@ -50,7 +51,7 @@ def run_pass(
 
 
 # We dont allow [cf, aod, schedule] appear in move function
-@ir.dialect_group(structural.union([action, spec, grid]))
+@ir.dialect_group(structural.union([action, spec, grid, filled]))
 def tweezer(self):
     fold_pass = Fold(self)
     typeinfer_pass = TypeInfer(self)
@@ -91,7 +92,7 @@ def run_pass(
 
 # no action allow. can have cf, with addtional spec
 @ir.dialect_group(
-    structural.union([init, schedule, path, grid, spec, gate, op, measure])
+    structural.union([init, schedule, path, grid, filled, spec, gate, op, measure])
 )
 def move(self):
     schedule_to_path = ScheduleToPath(self)