diff --git a/src/pystencils/nbackend/exceptions.py b/src/pystencils/nbackend/exceptions.py
index 8e48653047d70ccac3fea1d490871ff861510e9b..6fe83b3f9dad9f7cfe07dd90b65a887962b2e251 100644
--- a/src/pystencils/nbackend/exceptions.py
+++ b/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
diff --git a/src/pystencils/nbackend/kernelcreation/context.py b/src/pystencils/nbackend/kernelcreation/context.py
index 8d14a5d4f82cc61ec9d0be4a4ec5130bdcfd8b74..40cd2448f27ad08a03b35e66221d4f30cb6dab73 100644
--- a/src/pystencils/nbackend/kernelcreation/context.py
+++ b/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:
diff --git a/src/pystencils/nbackend/kernelcreation/domain_kernels.py b/src/pystencils/nbackend/kernelcreation/domain_kernels.py
index 1ee1cfae95463b39261b75a0dc6557e9431fe203..18f157c09bcde5f28775c0d75dd4eef73cfd0d8d 100644
--- a/src/pystencils/nbackend/kernelcreation/domain_kernels.py
+++ b/src/pystencils/nbackend/kernelcreation/domain_kernels.py
@@ -1,47 +1,55 @@
-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
+ )
diff --git a/src/pystencils/nbackend/kernelcreation/iteration_space.py b/src/pystencils/nbackend/kernelcreation/iteration_space.py
index ade56a04ca1fb62d31512c8ab11fee6da857f9de..1cabc6736209398e7cd10551c0ae7f980caa0f60 100644
--- a/src/pystencils/nbackend/kernelcreation/iteration_space.py
+++ b/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
)
]
diff --git a/src/pystencils/nbackend/kernelcreation/options.py b/src/pystencils/nbackend/kernelcreation/options.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9efd83c489eb9627bbf8883dcc1cd59e18adbca
--- /dev/null
+++ b/src/pystencils/nbackend/kernelcreation/options.py
@@ -0,0 +1,54 @@
+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."
+ )