diff --git a/src/pystencils/nbackend/arrays.py b/src/pystencils/nbackend/arrays.py
index 369f3d19af31fa471851fc9c9be20836034a03f6..3027c0b4096ab0a6bd8578613f2f5acd4e53d1bb 100644
--- a/src/pystencils/nbackend/arrays.py
+++ b/src/pystencils/nbackend/arrays.py
@@ -104,9 +104,15 @@ class PsLinearizedArray:
for i, s in enumerate(strides)
)
+ self._base_ptr = PsArrayBasePointer(f"{self._name}_data", self)
+
@property
def name(self):
return self._name
+
+ @property
+ def base_pointer(self) -> PsArrayBasePointer:
+ return self._base_ptr
@property
def shape(self) -> tuple[PsArrayShapeVar | PsTypedConstant, ...]:
diff --git a/src/pystencils/nbackend/ast/kernelfunction.py b/src/pystencils/nbackend/ast/kernelfunction.py
index 9aecb16eac40cbcbd84e3fbfab24fcd8ae923a39..b911f110091e251990dab4f5eb50d64733360545 100644
--- a/src/pystencils/nbackend/ast/kernelfunction.py
+++ b/src/pystencils/nbackend/ast/kernelfunction.py
@@ -67,7 +67,7 @@ class PsKernelFunction(PsAstNode):
__match_args__ = ("body",)
def __init__(self, body: PsBlock, target: Target, name: str = "kernel"):
- self._body = body
+ self._body: PsBlock = body
self._target = target
self._name = name
diff --git a/src/pystencils/nbackend/kernelcreation/context.py b/src/pystencils/nbackend/kernelcreation/context.py
index 154e06aefa1aee5caf5027e52f1f08bd0bcf446c..28417d56407cb8764e1bc2d3023db56ced113f2e 100644
--- a/src/pystencils/nbackend/kernelcreation/context.py
+++ b/src/pystencils/nbackend/kernelcreation/context.py
@@ -7,18 +7,18 @@ from abc import ABC
from ...field import Field
from ...typing import TypedSymbol, BasicType
-from ..arrays import PsLinearizedArray, PsArrayBasePointer
+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
@dataclass
-class PsFieldArrayPair:
+class PsArrayDescriptor:
field: Field
array: PsLinearizedArray
- base_ptr: PsArrayBasePointer
class IterationSpace(ABC):
@@ -80,47 +80,35 @@ class SparseIterationSpace(IterationSpace):
class KernelCreationContext:
- """Manages the translation process from the SymPy frontend to the backend AST.
+ """Manages the translation process from the SymPy frontend to the backend AST, and collects
+ all necessary information for the translation.
- It does the following things:
-
- - Default data types: The context knows the data types that should be applied by default
- to SymPy expressions.
- - Management of fields. The context manages all mappings from front-end `Field`s to their
- underlying `PsLinearizedArray`s.
- - Collection of constraints. All constraints that arise during translation are collected in the
- context, and finally attached to the kernel function object once translation is complete.
Data Types
----------
- - The `index_dtype` is the data type used throughout translation for all loop counters and array indexing.
- - The `default_numeric_dtype` is the data type assigned by default to all symbols occuring in SymPy assignments
+ The kernel creation context manages the default data types for loop limits and counters, index calculations,
+ and the typifier.
Fields and Arrays
- -----------------
-
- There's several types of fields that need to be mapped to arrays.
+ ------------------
- - `FieldType.GENERIC` corresponds to domain fields.
- Domain fields can only be accessed by relative offsets, and therefore must always
- be associated with an iteration space that provides a spatial index tuple.
- - `FieldType.INDEXED` are 1D arrays of index structures. They must be accessed by a single running index.
- If there is at least one indexed field present there must also exist an index source for that field
- (loop or device indexing).
- An indexed field may itself be an index source for domain fields.
- - `FieldType.BUFFER` are 1D arrays whose indices must be incremented with each access.
- Within a domain, a buffer may be either written to or read from, never both.
+ The kernel creation context acts as a factory for mapping fields to arrays.
+ Iteration Space
+ ---------------
- In the translator, frontend fields and backend arrays are managed together using the `PsFieldArrayPair` class.
+ The context manages the iteration space within which the current translation takes place. It may be a sparse
+ or full iteration space.
"""
def __init__(self, index_dtype: PsIntegerType):
self._index_dtype = index_dtype
- self._arrays: dict[Field, PsFieldArrayPair] = dict()
+ self._arrays: dict[Field, PsLinearizedArray] = dict()
self._constraints: list[PsKernelConstraint] = []
+ self._ispace: IterationSpace | None = None
+
@property
def index_dtype(self) -> PsIntegerType:
return self._index_dtype
@@ -132,35 +120,47 @@ class KernelCreationContext:
def constraints(self) -> tuple[PsKernelConstraint, ...]:
return tuple(self._constraints)
- def add_field(self, field: Field) -> PsFieldArrayPair:
- arr_shape = tuple(
- (
- Ellipsis if isinstance(s, TypedSymbol) else s
- ) # TODO: Field should also use ellipsis
- for s in field.shape
- )
+ def get_array(self, field: Field) -> PsLinearizedArray:
+ if field not in self._arrays:
+ arr_shape = tuple(
+ (
+ Ellipsis if isinstance(s, TypedSymbol) else s
+ ) # TODO: Field should also use ellipsis
+ for s in field.shape
+ )
+
+ arr_strides = tuple(
+ (
+ Ellipsis if isinstance(s, TypedSymbol) else s
+ ) # TODO: Field should also use ellipsis
+ for s in field.strides
+ )
- arr_strides = tuple(
- (
- Ellipsis if isinstance(s, TypedSymbol) else s
- ) # TODO: Field should also use ellipsis
- for s in field.strides
- )
+ # TODO: frontend should use new type system
+ element_type = make_type(cast(BasicType, field.dtype).numpy_dtype.type)
- # TODO: frontend should use new type system
- element_type = make_type(cast(BasicType, field.dtype).numpy_dtype.type)
+ arr = PsLinearizedArray(
+ field.name, element_type, arr_shape, arr_strides, self.index_dtype
+ )
- arr = PsLinearizedArray(
- field.name, element_type, arr_shape, arr_strides, self.index_dtype
- )
+ self._arrays[field] = arr
- fa_pair = PsFieldArrayPair(
- field=field, array=arr, base_ptr=PsArrayBasePointer("arr_data", arr)
- )
+ return self._arrays[field]
- self._arrays[field] = fa_pair
+ def set_iteration_space(self, ispace: IterationSpace):
+ self._ispace = ispace
- return fa_pair
+ def get_iteration_space(self) -> IterationSpace:
+ if self._ispace is None:
+ raise PsInternalCompilerError("No iteration space set in context.")
+ return self._ispace
- def get_array_descriptor(self, field: Field):
- return self._arrays[field]
+ def get_full_iteration_space(self) -> FullIterationSpace:
+ if not isinstance(self._ispace, FullIterationSpace):
+ raise PsInternalCompilerError("No full iteration space set in context.")
+ return self._ispace
+
+ def get_sparse_iteration_space(self) -> SparseIterationSpace:
+ if not isinstance(self._ispace, SparseIterationSpace):
+ raise PsInternalCompilerError("No sparse iteration space set in context.")
+ return self._ispace
diff --git a/src/pystencils/nbackend/kernelcreation/defaults.py b/src/pystencils/nbackend/kernelcreation/defaults.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ef68fe89fe989b01cdc699870b63fe06b05db9f
--- /dev/null
+++ b/src/pystencils/nbackend/kernelcreation/defaults.py
@@ -0,0 +1,17 @@
+"""This module defines various default types, symbols and variables for use in pystencils kernels.
+
+On many occasions the SymPy frontend uses canonical symbols and types.
+With the pymbolic-based backend, these symbols have to exist in two
+variants; as `sp.Symbol` or `TypedSymbol`, and as `PsTypedVariable`s.
+Therefore, for conciseness, this module should collect and provide each of these symbols.
+
+We might furthermore consider making the defaults collection configurable.
+
+A possibly incomplete list of symbols and types that need to be defined:
+
+ - The default indexing data type (currently loosely defined as `int`)
+ - The default spatial iteration counters (currently defined by `LoopOverCoordinate`)
+ - The names of the coordinate members of index lists (currently in `CreateKernelConfig.coordinate_names`)
+ - The sparse iteration counter (doesn't even exist yet)
+ - ...
+"""
diff --git a/src/pystencils/nbackend/kernelcreation/domain_kernels.py b/src/pystencils/nbackend/kernelcreation/domain_kernels.py
index d76ebbf2bc37958f60a18675cfdc20dac12c0eac..a66cbd12ef05e57ffc972910045aff964d7fb2a1 100644
--- a/src/pystencils/nbackend/kernelcreation/domain_kernels.py
+++ b/src/pystencils/nbackend/kernelcreation/domain_kernels.py
@@ -21,12 +21,6 @@ def create_domain_kernel(assignments: AssignmentCollection):
check = KernelConstraintsCheck() # TODO: config
check.visit(assignments)
- all_fields: set[Field] = check.fields_written | check.fields_read
-
- # 3. Register fields
- for f in all_fields:
- ctx.add_field(f)
-
# All steps up to this point are the same in domain and indexed kernels;
# the difference now comes with the iteration space.
#
@@ -35,7 +29,7 @@ def create_domain_kernel(assignments: AssignmentCollection):
# Indexed kernels, on the other hand, have to create a sparse iteration space
# from one index list.
- # 4. Create iteration space
+ # 3. Create iteration space
ghost_layers: int = NotImplemented # determine required ghost layers
common_shape: tuple[
int | EllipsisType, ...
@@ -45,12 +39,12 @@ def create_domain_kernel(assignments: AssignmentCollection):
NotImplemented # create from ghost layers and with given shape
)
- # 5. Freeze assignments
+ # 4. Freeze assignments
# This call is the same for both domain and indexed kernels
- freeze = FreezeExpressions(ctx, ispace)
+ freeze = FreezeExpressions(ctx)
kernel_body: PsBlock = freeze(assignments)
- # 6. Typify
+ # 5. Typify
# Also the same for both types of kernels
# determine_types(kernel_body)
@@ -58,13 +52,13 @@ def create_domain_kernel(assignments: AssignmentCollection):
# Here we might hand off the compilation to a target-specific part of the compiler
# (CPU/CUDA/...), since these will likely also apply very different optimizations.
- # 7. Add loops or device indexing
+ # 6. Add loops or device indexing
# This step translates the iteration space to actual index calculation code and is once again
# different in indexed and domain kernels.
- # 8. Apply optimizations
+ # 7. Apply optimizations
# - Vectorization
# - OpenMP
# - Loop Splitting, Tiling, Blocking
- # 9. Create and return kernel function.
+ # 8. Create and return kernel function.
diff --git a/src/pystencils/nbackend/kernelcreation/freeze.py b/src/pystencils/nbackend/kernelcreation/freeze.py
index 7ae3879934366309ae8aaf5d11f5711114a13804..e9fc6ccd72bab3f9dac24f3bcd7539c9a3923d5b 100644
--- a/src/pystencils/nbackend/kernelcreation/freeze.py
+++ b/src/pystencils/nbackend/kernelcreation/freeze.py
@@ -3,19 +3,17 @@ from pymbolic.interop.sympy import SympyToPymbolicMapper
from ...field import Field, FieldType
-from .context import KernelCreationContext, IterationSpace, SparseIterationSpace
+from .context import KernelCreationContext
from ..ast.nodes import PsAssignment
from ..types import PsSignedIntegerType, PsIeeeFloatType, PsUnsignedIntegerType
from ..typed_expressions import PsTypedVariable
from ..arrays import PsArrayAccess
-from ..exceptions import PsInternalCompilerError
class FreezeExpressions(SympyToPymbolicMapper):
- def __init__(self, ctx: KernelCreationContext, ispace: IterationSpace):
+ def __init__(self, ctx: KernelCreationContext):
self._ctx = ctx
- self._ispace = ispace
def map_Assignment(self, expr): # noqa
lhs = self.rec(expr.lhs)
@@ -44,9 +42,8 @@ class FreezeExpressions(SympyToPymbolicMapper):
def map_Access(self, access: Field.Access):
field = access.field
- array_desc = self._ctx.get_array_descriptor(field)
- array = array_desc.array
- ptr = array_desc.base_ptr
+ array = self._ctx.get_array(field)
+ ptr = array.base_pointer
offsets: list[pb.Expression] = [self.rec(o) for o in access.offsets]
indices: list[pb.Expression] = [self.rec(o) for o in access.index]
@@ -56,17 +53,14 @@ class FreezeExpressions(SympyToPymbolicMapper):
case FieldType.GENERIC:
# Add the iteration counters
offsets = [
- i + o for i, o in zip(self._ispace.spatial_indices, offsets)
+ i + o for i, o in zip(self._ctx.get_iteration_space().spatial_indices, offsets)
]
case FieldType.INDEXED:
- if isinstance(self._ispace, SparseIterationSpace):
- # TODO: make sure index (and all offsets?) are zero
- # TODO: Add sparse iteration counter
- raise NotImplementedError()
- else:
- raise PsInternalCompilerError(
- "Cannot translate index field access without a sparse iteration space."
- )
+ # flake8: noqa
+ sparse_ispace = self._ctx.get_sparse_iteration_space()
+ # TODO: make sure index (and all offsets?) are zero
+ # TODO: Add sparse iteration counter
+ raise NotImplementedError()
case FieldType.CUSTOM:
raise ValueError("Custom fields support only absolute accesses.")
case unknown: