extend iteration space; create it.space from ghost layers

src/pystencils/nbackend/kernelcreation/context.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 from typing import cast
 from dataclasses import dataclass
 
-from abc import ABC
 
 from ...field import Field
 from ...typing import TypedSymbol, BasicType
@@ -10,10 +9,11 @@ from ...typing import TypedSymbol, BasicType
 from ..arrays import PsLinearizedArray
 from ..types import PsIntegerType
 from ..types.quick import make_type
-from ..typed_expressions import PsTypedVariable, VarOrConstant
 from ..constraints import PsKernelConstraint
 from ..exceptions import PsInternalCompilerError
 
+from .iteration_space import IterationSpace, FullIterationSpace, SparseIterationSpace
+
 
 @dataclass
 class PsArrayDescriptor:
@@ -21,64 +21,6 @@ class PsArrayDescriptor:
     array: PsLinearizedArray
 
 
-class IterationSpace(ABC):
-    """Represents the n-dimensonal iteration space of a pystencils kernel.
-
-    Instances of this class represent the kernel's iteration region during translation from
-    SymPy, before any indexing sources are generated. It provides the counter symbols which
-    should be used to translate field accesses to array accesses.
-
-    There are two types of iteration spaces, modelled by subclasses:
-     - The full iteration space translates to an n-dimensional loop nest or the corresponding device
-       indexing scheme.
-     - The sparse iteration space translates to a single loop over an index list which in turn provides the
-       spatial indices.
-    """
-
-    def __init__(self, spatial_indices: tuple[PsTypedVariable, ...]):
-        if len(spatial_indices) == 0:
-            raise ValueError("Iteration space must be at least one-dimensional.")
-
-        self._spatial_indices = spatial_indices
-
-    @property
-    def spatial_indices(self) -> tuple[PsTypedVariable, ...]:
-        return self._spatial_indices
-
-
-class FullIterationSpace(IterationSpace):
-    def __init__(
-        self,
-        lower: tuple[VarOrConstant, ...],
-        upper: tuple[VarOrConstant, ...],
-        counters: tuple[PsTypedVariable, ...],
-    ):
-        if not (len(lower) == len(upper) == len(counters)):
-            raise ValueError(
-                "Lower and upper iteration limits and counters must have the same shape."
-            )
-
-        super().__init__(counters)
-
-        self._lower = lower
-        self._upper = upper
-        self._counters = counters
-
-    @property
-    def lower(self):
-        return self._lower
-
-    @property
-    def upper(self):
-        return self._upper
-
-
-class SparseIterationSpace(IterationSpace):
-    def __init__(self, spatial_index_variables: tuple[PsTypedVariable, ...]):
-        super().__init__(spatial_index_variables)
-        # todo
-
-
 class KernelCreationContext:
     """Manages the translation process from the SymPy frontend to the backend AST, and collects
     all necessary information for the translation.

src/pystencils/nbackend/kernelcreation/defaults.py

@@ -15,3 +15,41 @@ A possibly incomplete list of symbols and types that need to be defined:
  - The sparse iteration counter (doesn't even exist yet)
  - ...
 """
+
+from typing import TypeVar, Generic, Callable
+from ..types import PsAbstractType, PsSignedIntegerType
+from ..typed_expressions import PsTypedVariable
+
+from ...typing import TypedSymbol
+
+SymbolT = TypeVar("SymbolT")
+
+
+class PsDefaults(Generic[SymbolT]):
+    def __init__(self, symcreate: Callable[[str, PsAbstractType], SymbolT]):
+        self.index_dtype = PsSignedIntegerType(64)
+        """Default data type for indices."""
+
+        self.spatial_counter_names = ("ctr_0", "ctr_1", "ctr_2")
+        """Names of the default spatial counters"""
+
+        self.spatial_counters = (
+            symcreate("ctr_0", self.index_dtype),
+            symcreate("ctr_1", self.index_dtype),
+            symcreate("ctr_2", self.index_dtype),
+        )
+        """Default spatial counters"""
+
+        self.index_struct_coordinates = (
+            symcreate("x", self.index_dtype),
+            symcreate("y", self.index_dtype),
+            symcreate("z", self.index_dtype),
+        )
+        """Default symbols for spatial coordinates in index list structures"""
+
+        self.sparse_iteration_counter = symcreate("list_idx", self.index_dtype)
+        """Default sparse iteration counter."""
+
+
+Sympy = PsDefaults[TypedSymbol](TypedSymbol)
+Pymbolic = PsDefaults[PsTypedVariable](PsTypedVariable)

src/pystencils/nbackend/kernelcreation/iteration_space.py

+from __future__ import annotations
+from typing import Sequence, TYPE_CHECKING
+from abc import ABC
+from dataclasses import dataclass
+from functools import reduce
+from operator import mul
+
+from ...field import Field
+from ...transformations import get_common_field
+
+from ..typed_expressions import (
+    PsTypedVariable,
+    VarOrConstant,
+    ExprOrConstant,
+    PsTypedConstant,
+)
+from ..arrays import PsLinearizedArray
+from .defaults import Pymbolic as Defaults
+
+if TYPE_CHECKING:
+    from .context import KernelCreationContext
+
+
+class IterationSpace(ABC):
+    """Represents the n-dimensonal iteration space of a pystencils kernel.
+
+    Instances of this class represent the kernel's iteration region during translation from
+    SymPy, before any indexing sources are generated. It provides the counter symbols which
+    should be used to translate field accesses to array accesses.
+
+    There are two types of iteration spaces, modelled by subclasses:
+     - The full iteration space translates to an n-dimensional loop nest or the corresponding device
+       indexing scheme.
+     - The sparse iteration space translates to a single loop over an index list which in turn provides the
+       spatial indices.
+    """
+
+    def __init__(self, spatial_indices: tuple[PsTypedVariable, ...]):
+        if len(spatial_indices) == 0:
+            raise ValueError("Iteration space must be at least one-dimensional.")
+
+        self._spatial_indices = spatial_indices
+
+    @property
+    def spatial_indices(self) -> tuple[PsTypedVariable, ...]:
+        return self._spatial_indices
+
+
+class FullIterationSpace(IterationSpace):
+    @dataclass
+    class Dimension:
+        start: VarOrConstant
+        stop: VarOrConstant
+        step: VarOrConstant
+        counter: PsTypedVariable
+
+    @staticmethod
+    def create_with_ghost_layers(
+        ctx: KernelCreationContext,
+        fields: set[Field],
+        ghost_layers: int | Sequence[int | tuple[int, int]],
+    ) -> FullIterationSpace:
+        """Create an iteration space for a collection of fields with ghost layers."""
+
+        repr_field: Field = get_common_field(fields)
+        repr_arr = ctx.get_array(repr_field)
+        dim = repr_field.spatial_dimensions
+        counters = [
+            PsTypedVariable(name, ctx.index_dtype)
+            for name in Defaults.spatial_counter_names
+        ]
+
+        if isinstance(ghost_layers, int):
+            ghost_layers_spec = [(ghost_layers, ghost_layers) for _ in range(dim)]
+        else:
+            if len(ghost_layers) != dim:
+                raise ValueError("Too few entries in ghost layer spec")
+            ghost_layers_spec = [
+                ((gl, gl) if isinstance(gl, int) else gl) for gl in ghost_layers
+            ]
+
+        one = PsTypedConstant(1, ctx.index_dtype)
+
+        ghost_layer_exprs = [
+            (
+                PsTypedConstant(gl_left, ctx.index_dtype),
+                PsTypedConstant(gl_right, ctx.index_dtype),
+            )
+            for (gl_left, gl_right) in ghost_layers_spec
+        ]
+
+        dimensions = [
+            FullIterationSpace.Dimension(gl_left, shape - gl_right, one, ctr)
+            for (gl_left, gl_right), shape, ctr in zip(
+                ghost_layer_exprs, repr_arr.shape, counters, strict=True
+            )
+        ]
+
+        return FullIterationSpace(ctx, dimensions)
+
+    def __init__(self, ctx: KernelCreationContext, dimensions: Sequence[Dimension]):
+        super().__init__(tuple(dim.counter for dim in dimensions))
+
+        self._ctx = ctx
+        self._dimensions = dimensions
+
+    @property
+    def dimensions(self):
+        return self._dimensions
+
+    @property
+    def lower(self):
+        return (dim.start for dim in self._dimensions)
+
+    @property
+    def upper(self):
+        return (dim.stop for dim in self._dimensions)
+
+    @property
+    def steps(self):
+        return (dim.step for dim in self._dimensions)
+
+    def num_iteration_items(self, dimension: int | None = None) -> ExprOrConstant:
+        if dimension is None:
+            return reduce(
+                mul, (self.num_iteration_items(d) for d in range(len(self.dimensions)))
+            )
+        else:
+            dim = self.dimensions[dimension]
+            one = PsTypedConstant(1, self._ctx.index_dtype)
+            return one + (dim.stop - dim.start - one) / dim.step
+
+
+class SparseIterationSpace(IterationSpace):
+    def __init__(
+        self,
+        spatial_indices: tuple[PsTypedVariable, ...],
+        index_list: PsLinearizedArray,
+    ):
+        super().__init__(spatial_indices)
+        self._index_list = index_list
+
+    @property
+    def index_list(self) -> PsLinearizedArray:
+        return self._index_list