From 9d95a9593763ed40783b751d9ce5a9da86da7f66 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Fri, 19 Jan 2024 23:40:50 +0100
Subject: [PATCH] extend iteration space; create it.space from ghost layers
---
.../nbackend/kernelcreation/context.py | 62 +-------
.../nbackend/kernelcreation/defaults.py | 38 +++++
.../kernelcreation/iteration_space.py | 145 ++++++++++++++++++
3 files changed, 185 insertions(+), 60 deletions(-)
create mode 100644 src/pystencils/nbackend/kernelcreation/iteration_space.py
diff --git a/src/pystencils/nbackend/kernelcreation/context.py b/src/pystencils/nbackend/kernelcreation/context.py
index 28417d564..8d14a5d4f 100644
--- a/src/pystencils/nbackend/kernelcreation/context.py
+++ b/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.
diff --git a/src/pystencils/nbackend/kernelcreation/defaults.py b/src/pystencils/nbackend/kernelcreation/defaults.py
index 2ef68fe89..c928ef4e4 100644
--- a/src/pystencils/nbackend/kernelcreation/defaults.py
+++ b/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)
diff --git a/src/pystencils/nbackend/kernelcreation/iteration_space.py b/src/pystencils/nbackend/kernelcreation/iteration_space.py
new file mode 100644
index 000000000..ade56a04c
--- /dev/null
+++ b/src/pystencils/nbackend/kernelcreation/iteration_space.py
@@ -0,0 +1,145 @@
+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
--
GitLab