create full iteration space

src/pystencils/nbackend/exceptions.py

@@ -2,5 +2,13 @@ class PsInternalCompilerError(Exception):
     pass
 
 
+class PsOptionsError(Exception):
+    pass
+
+
+class PsInputError(Exception):
+    pass
+
+
 class PsMalformedAstException(Exception):
     pass

src/pystencils/nbackend/kernelcreation/context.py

@@ -3,7 +3,7 @@ from typing import cast
 from dataclasses import dataclass
 
 
-from ...field import Field
+from ...field import Field, FieldType
 from ...typing import TypedSymbol, BasicType
 
 from ..arrays import PsLinearizedArray
@@ -12,13 +12,16 @@ from ..types.quick import make_type
 from ..constraints import PsKernelConstraint
 from ..exceptions import PsInternalCompilerError
 
+from .options import KernelCreationOptions
 from .iteration_space import IterationSpace, FullIterationSpace, SparseIterationSpace
 
 
 @dataclass
-class PsArrayDescriptor:
-    field: Field
-    array: PsLinearizedArray
+class FieldsInKernel:
+    domain_fields: set[Field] = set()
+    index_fields: set[Field] = set()
+    custom_fields: set[Field] = set()
+    buffer_fields: set[Field] = set()
 
 
 class KernelCreationContext:
@@ -44,16 +47,21 @@ class KernelCreationContext:
     or full iteration space.
     """
 
-    def __init__(self, index_dtype: PsIntegerType):
-        self._index_dtype = index_dtype
+    def __init__(self, options: KernelCreationOptions):
+        self._options = options
         self._arrays: dict[Field, PsLinearizedArray] = dict()
         self._constraints: list[PsKernelConstraint] = []
 
+        self._fields_collection = FieldsInKernel()
         self._ispace: IterationSpace | None = None
 
+    @property
+    def options(self) -> KernelCreationOptions:
+        return self._options
+
     @property
     def index_dtype(self) -> PsIntegerType:
-        return self._index_dtype
+        return self._options.index_dtype
 
     def add_constraints(self, *constraints: PsKernelConstraint):
         self._constraints += constraints
@@ -62,6 +70,24 @@ class KernelCreationContext:
     def constraints(self) -> tuple[PsKernelConstraint, ...]:
         return tuple(self._constraints)
 
+    @property
+    def fields(self) -> FieldsInKernel:
+        return self._fields_collection
+
+    def add_field(self, field: Field):
+        """Add the given field to the context's fields collection"""
+        match field.field_type:
+            case FieldType.GENERIC | FieldType.STAGGERED | FieldType.STAGGERED_FLUX:
+                self._fields_collection.domain_fields.add(field)
+            case FieldType.BUFFER:
+                self._fields_collection.buffer_fields.add(field)
+            case FieldType.INDEXED:
+                self._fields_collection.index_fields.add(field)
+            case FieldType.CUSTOM:
+                self._fields_collection.custom_fields.add(field)
+            case _:
+                assert False, "unreachable code"
+
     def get_array(self, field: Field) -> PsLinearizedArray:
         if field not in self._arrays:
             arr_shape = tuple(
@@ -90,6 +116,8 @@ class KernelCreationContext:
         return self._arrays[field]
 
     def set_iteration_space(self, ispace: IterationSpace):
+        if self._ispace is not None:
+            raise PsInternalCompilerError("Iteration space was already set.")
         self._ispace = ispace
 
     def get_iteration_space(self) -> IterationSpace:

src/pystencils/nbackend/kernelcreation/domain_kernels.py

-from types import EllipsisType
+from typing import Sequence
+from itertools import chain
 
 from ...simp import AssignmentCollection
-from ...field import Field
+from ...field import Field, FieldType
 from ...kernel_contrains_check import KernelConstraintsCheck
 
-from ..types.quick import SInt
 from ..ast import PsBlock
 
-from .context import KernelCreationContext, FullIterationSpace
+from .context import KernelCreationContext, IterationSpace
 from .freeze import FreezeExpressions
 from .typification import Typifier
+from .options import KernelCreationOptions
+from ..exceptions import PsInputError, PsInternalCompilerError
 
 # flake8: noqa
 
 
-def create_domain_kernel(assignments: AssignmentCollection):
-    #   TODO: Assemble configuration
-
+def create_kernel(assignments: AssignmentCollection, options: KernelCreationOptions):
     #   1. Prepare context
-    ctx = KernelCreationContext(SInt(64))  # TODO: how to determine index type?
+    ctx = KernelCreationContext(options)
 
     #   2. Check kernel constraints and collect all fields
+
+    """
+    TODO: Replace the KernelConstraintsCheck by a KernelAnalysis pass.
+
+    The kernel analysis should:
+     - Check constraints on the assignments (SSA form, independence conditions, ...)
+     - Collect all fields and register them in the context
+     - Maybe collect all field accesses and register them at the context
+     
+    Running all this analysis in a single pass will likely improve compiler performance
+    since no additional searches, e.g. for field accesses, are necessary later.
+    """
+
     check = KernelConstraintsCheck()  # TODO: config
     check.visit(assignments)
 
-    #   All steps up to this point are the same in domain and indexed kernels;
-    #   the difference now comes with the iteration space.
-    #
-    #   Domain kernels create a full iteration space from their iteration slice
-    #   which is either explicitly given or computed from ghost layer requirements.
-    #   Indexed kernels, on the other hand, have to create a sparse iteration space
-    #   from one index list.
+    #   Collect all fields
+    for f in chain(check.fields_written, check.fields_read):
+        ctx.add_field(f)
 
     #   3. Create iteration space
-    ghost_layers: int = NotImplemented  # determine required ghost layers
-    common_shape: tuple[
-        int | EllipsisType, ...
-    ] = NotImplemented  # unify field shapes, add parameter constraints
-    #   don't forget custom iteration slice
-    ispace: FullIterationSpace = (
-        NotImplemented  # create from ghost layers and with given shape
+    ispace: IterationSpace = (
+        create_sparse_iteration_space(ctx, assignments)
+        if len(ctx.fields.index_fields) > 0
+        else create_full_iteration_space(ctx, assignments)
     )
 
+    ctx.set_iteration_space(ispace)
+
     #   4. Freeze assignments
     #   This call is the same for both domain and indexed kernels
     freeze = FreezeExpressions(ctx)
@@ -66,3 +74,63 @@ def create_domain_kernel(assignments: AssignmentCollection):
     #     - Loop Splitting, Tiling, Blocking
 
     #   8. Create and return kernel function.
+
+
+def create_sparse_iteration_space(
+    ctx: KernelCreationContext, assignments: AssignmentCollection
+) -> IterationSpace:
+    return NotImplemented
+
+
+def create_full_iteration_space(
+    ctx: KernelCreationContext, assignments: AssignmentCollection
+) -> IterationSpace:
+    assert not ctx.fields.index_fields
+
+    #   Collect all relative accesses into domain fields
+    def access_filter(acc: Field.Access):
+        return acc.field.field_type in (
+            FieldType.GENERIC,
+            FieldType.STAGGERED,
+            FieldType.STAGGERED_FLUX,
+        )
+
+    domain_field_accesses = assignments.atoms(Field.Access)
+    domain_field_accesses = set(filter(access_filter, domain_field_accesses))
+
+    # The following scenarios exist:
+    # - We have at least one domain field -> find the common field and use it to determine the iteration region
+    # - We have no domain fields, but at least one custom field -> determine common field from custom fields
+    # - We have neither domain nor custom fields -> Error
+
+    from ...transformations import get_common_field
+
+    if len(domain_field_accesses) > 0:
+        archetype_field = get_common_field(ctx.fields.domain_fields)
+        inferred_gls = max(
+            [fa.required_ghost_layers for fa in domain_field_accesses]
+        )
+    elif len(ctx.fields.custom_fields) > 0:
+        archetype_field = get_common_field(ctx.fields.custom_fields)
+        inferred_gls = 0
+    else:
+        raise PsInputError(
+            "Unable to construct iteration space: The kernel contains no accesses to domain or custom fields."
+        )
+
+    # If the user provided a ghost layer specification, use that
+    # Otherwise, if an iteration slice was specified, use that
+    # Otherwise, use the inferred ghost layers
+
+    from .iteration_space import FullIterationSpace
+
+    if ctx.options.ghost_layers is not None:
+        return FullIterationSpace.create_with_ghost_layers(
+            ctx, archetype_field, ctx.options.ghost_layers
+        )
+    elif ctx.options.iteration_slice is not None:
+        raise PsInternalCompilerError("Iteration slices not supported yet")
+    else:
+        return FullIterationSpace.create_with_ghost_layers(
+            ctx, archetype_field, inferred_gls
+        )

src/pystencils/nbackend/kernelcreation/iteration_space.py

@@ -6,7 +6,6 @@ from functools import reduce
 from operator import mul
 
 from ...field import Field
-from ...transformations import get_common_field
 
 from ..typed_expressions import (
     PsTypedVariable,
@@ -57,14 +56,13 @@ class FullIterationSpace(IterationSpace):
     @staticmethod
     def create_with_ghost_layers(
         ctx: KernelCreationContext,
-        fields: set[Field],
+        archetype_field: 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
+        archetype_array = ctx.get_array(archetype_field)
+        dim = archetype_field.spatial_dimensions
         counters = [
             PsTypedVariable(name, ctx.index_dtype)
             for name in Defaults.spatial_counter_names
@@ -92,7 +90,7 @@ class FullIterationSpace(IterationSpace):
         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
+                ghost_layer_exprs, archetype_array.shape, counters, strict=True
             )
         ]
 

src/pystencils/nbackend/kernelcreation/options.py

+from typing import Sequence
+from dataclasses import dataclass
+
+from ...enums import Target
+from ..exceptions import PsOptionsError
+from ..types import PsIntegerType
+
+from .defaults import Sympy as SpDefaults
+
+
+@dataclass
+class KernelCreationOptions:
+    """Options for create_kernel."""
+
+    target: Target = Target.CPU
+    """The code generation target.
+    
+    TODO: Enhance `Target` from enum to a larger target spec, e.g. including vectorization architecture, ...
+    """
+
+    function_name: str = "kernel"
+    """Name of the generated function"""
+
+    ghost_layers: None | int | Sequence[int | tuple[int, int]] = None
+    """Specifies the number of ghost layers of the iteration region.
+    
+    Options:
+     - `None`: Required ghost layers are inferred from field accesses
+     - `int`:  A uniform number of ghost layers in each spatial coordinate is applied
+     - `Sequence[int, tuple[int, int]]`: Ghost layers are specified for each spatial coordinate.
+        In each coordinate, a single integer specifies the ghost layers at both the lower and upper iteration limit,
+        while a pair of integers specifies the lower and upper ghost layers separately.
+
+    When manually specifying ghost layers, it is the user's responsibility to avoid out-of-bounds memory accesses.
+    If `ghost_layers=None` is specified, the iteration region may otherwise be set using the `iteration_slice` option.
+    """
+
+    iteration_slice: None | tuple[slice, ...] = None
+    """Specifies the kernel's iteration slice.
+    
+    `iteration_slice` may only be set if `ghost_layers = None`.
+    If it is set, a slice must be specified for each spatial coordinate.
+    TODO: Specification of valid slices and their behaviour
+    """
+
+    index_dtype: PsIntegerType = SpDefaults.index_dtype
+    """Data type used for all index calculations."""
+
+    def __post_init__(self):
+        if self.iteration_slice is not None and self.ghost_layers is not None:
+            raise PsOptionsError(
+                "Parameters `iteration_slice` and `ghost_layers` are mutually exclusive; "
+                "at most one of them may be set."
+            )