more simplifications, cleanup, and notes

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, ...]:

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
 

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

src/pystencils/nbackend/kernelcreation/defaults.py

+"""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)
+ - ...
+"""

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.

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: