Skip to content
Snippets Groups Projects
Commit 8490e6cb authored by Frederik Hennig's avatar Frederik Hennig
Browse files

move CreateKernelConfig to kernelcreation.py

parent 38643d15
No related branches found
No related tags found
No related merge requests found
Pipeline #62655 failed
from typing import Sequence
from dataclasses import dataclass
from ...enums import Target
from ...field import Field, FieldType
from ..jit import JitBase
from ..exceptions import PsOptionsError
from ..types import PsIntegerType, PsNumericType, PsIeeeFloatType
from .defaults import Sympy as SpDefaults
@dataclass
class CreateKernelConfig:
"""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, ...
"""
jit: JitBase | None = None
"""Just-in-time compiler used to compile and load the kernel for invocation from the current Python environment.
If left at `None`, a default just-in-time compiler will be inferred from the `target` parameter.
To explicitly disable JIT compilation, pass `nbackend.jit.no_jit`.
"""
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_field: Field | None = None
"""Index field for a sparse kernel.
If this option is set, a sparse kernel with the given field as index field will be generated.
"""
"""Data Types"""
index_dtype: PsIntegerType = SpDefaults.index_dtype
"""Data type used for all index calculations."""
default_dtype: PsNumericType = PsIeeeFloatType(64)
"""Default numeric data type.
This data type will be applied to all untyped symbols.
"""
def __post_init__(self):
# Check iteration space argument consistency
if (
int(self.iteration_slice is not None)
+ int(self.ghost_layers is not None)
+ int(self.index_field is not None)
> 1
):
raise PsOptionsError(
"Parameters `iteration_slice`, `ghost_layers` and 'index_field` are mutually exclusive; "
"at most one of them may be set."
)
# Check index field
if (
self.index_field is not None
and self.index_field.field_type != FieldType.INDEXED
):
raise PsOptionsError(
"Only fields with `field_type == FieldType.INDEXED` can be specified as `index_field`"
)
# Infer JIT
if self.jit is None:
match self.target:
case Target.CPU:
from ..jit import legacy_cpu
self.jit = legacy_cpu
case _:
raise NotImplementedError(f"No default JIT compiler implemented yet for target {self.target}")
...@@ -3,13 +3,12 @@ from __future__ import annotations ...@@ -3,13 +3,12 @@ from __future__ import annotations
from ...field import Field, FieldType from ...field import Field, FieldType
from ...sympyextensions.typed_sympy import TypedSymbol, BasicType, StructType from ...sympyextensions.typed_sympy import TypedSymbol, BasicType, StructType
from ..arrays import PsLinearizedArray from ..arrays import PsLinearizedArray
from ..types import PsIntegerType from ..types import PsIntegerType, PsNumericType
from ..types.quick import make_type from ..types.quick import make_type
from ..constraints import PsKernelConstraint from ..constraints import PsKernelConstraint
from ..exceptions import PsInternalCompilerError, KernelConstraintsError from ..exceptions import PsInternalCompilerError, KernelConstraintsError
from .defaults import Pymbolic as PbDefaults
from .config import CreateKernelConfig
from .iteration_space import IterationSpace, FullIterationSpace, SparseIterationSpace from .iteration_space import IterationSpace, FullIterationSpace, SparseIterationSpace
...@@ -44,8 +43,11 @@ class KernelCreationContext: ...@@ -44,8 +43,11 @@ class KernelCreationContext:
or full iteration space. or full iteration space.
""" """
def __init__(self, options: CreateKernelConfig): def __init__(self,
self._options = options default_dtype: PsNumericType = PbDefaults.numeric_dtype,
index_dtype: PsIntegerType = PbDefaults.index_dtype):
self._default_dtype = default_dtype
self._index_dtype = index_dtype
self._arrays: dict[Field, PsLinearizedArray] = dict() self._arrays: dict[Field, PsLinearizedArray] = dict()
self._constraints: list[PsKernelConstraint] = [] self._constraints: list[PsKernelConstraint] = []
...@@ -53,12 +55,12 @@ class KernelCreationContext: ...@@ -53,12 +55,12 @@ class KernelCreationContext:
self._ispace: IterationSpace | None = None self._ispace: IterationSpace | None = None
@property @property
def options(self) -> CreateKernelConfig: def default_dtype(self) -> PsNumericType:
return self._options return self._default_dtype
@property @property
def index_dtype(self) -> PsIntegerType: def index_dtype(self) -> PsIntegerType:
return self._options.index_dtype return self._index_dtype
def add_constraints(self, *constraints: PsKernelConstraint): def add_constraints(self, *constraints: PsKernelConstraint):
self._constraints += constraints self._constraints += constraints
......
...@@ -17,7 +17,7 @@ A possibly incomplete list of symbols and types that need to be defined: ...@@ -17,7 +17,7 @@ A possibly incomplete list of symbols and types that need to be defined:
""" """
from typing import TypeVar, Generic, Callable from typing import TypeVar, Generic, Callable
from ..types import PsAbstractType, PsSignedIntegerType, PsStructType from ..types import PsAbstractType, PsIeeeFloatType, PsSignedIntegerType, PsStructType
from ..typed_expressions import PsTypedVariable from ..typed_expressions import PsTypedVariable
from pystencils.sympyextensions.typed_sympy import TypedSymbol from pystencils.sympyextensions.typed_sympy import TypedSymbol
...@@ -27,6 +27,9 @@ SymbolT = TypeVar("SymbolT") ...@@ -27,6 +27,9 @@ SymbolT = TypeVar("SymbolT")
class PsDefaults(Generic[SymbolT]): class PsDefaults(Generic[SymbolT]):
def __init__(self, symcreate: Callable[[str, PsAbstractType], SymbolT]): def __init__(self, symcreate: Callable[[str, PsAbstractType], SymbolT]):
self.numeric_dtype = PsIeeeFloatType(64)
"""Default data type for numerical computations"""
self.index_dtype = PsSignedIntegerType(64) self.index_dtype = PsSignedIntegerType(64)
"""Default data type for indices.""" """Default data type for indices."""
......
...@@ -224,7 +224,7 @@ def get_archetype_field( ...@@ -224,7 +224,7 @@ def get_archetype_field(
def create_sparse_iteration_space( def create_sparse_iteration_space(
ctx: KernelCreationContext, assignments: AssignmentCollection ctx: KernelCreationContext, assignments: AssignmentCollection, index_field: Field | None = None
) -> IterationSpace: ) -> IterationSpace:
# All domain and custom fields must have the same spatial dimensions # All domain and custom fields must have the same spatial dimensions
# TODO: Must all domain fields have the same shape? # TODO: Must all domain fields have the same shape?
...@@ -242,9 +242,8 @@ def create_sparse_iteration_space( ...@@ -242,9 +242,8 @@ def create_sparse_iteration_space(
] ]
# Determine index field # Determine index field
if ctx.options.index_field is not None: if index_field is not None:
idx_field = ctx.options.index_field idx_arr = ctx.get_array(index_field)
idx_arr = ctx.get_array(idx_field)
idx_struct_type: PsStructType = failing_cast(PsStructType, idx_arr.element_type) idx_struct_type: PsStructType = failing_cast(PsStructType, idx_arr.element_type)
for coord in coord_members: for coord in coord_members:
...@@ -269,10 +268,16 @@ def create_sparse_iteration_space( ...@@ -269,10 +268,16 @@ def create_sparse_iteration_space(
def create_full_iteration_space( def create_full_iteration_space(
ctx: KernelCreationContext, assignments: AssignmentCollection ctx: KernelCreationContext,
assignments: AssignmentCollection,
ghost_layers: None | int | Sequence[int | tuple[int, int]] = None,
iteration_slice: None | tuple[slice, ...] = None
) -> IterationSpace: ) -> IterationSpace:
assert not ctx.fields.index_fields assert not ctx.fields.index_fields
if (ghost_layers is not None) and (iteration_slice is not None):
raise ValueError("At most one of `ghost_layers` and `iteration_slice` may be specified.")
# Collect all relative accesses into domain fields # Collect all relative accesses into domain fields
def access_filter(acc: Field.Access): def access_filter(acc: Field.Access):
return acc.field.field_type in ( return acc.field.field_type in (
...@@ -305,11 +310,11 @@ def create_full_iteration_space( ...@@ -305,11 +310,11 @@ def create_full_iteration_space(
# Otherwise, if an iteration slice was specified, use that # Otherwise, if an iteration slice was specified, use that
# Otherwise, use the inferred ghost layers # Otherwise, use the inferred ghost layers
if ctx.options.ghost_layers is not None: if ghost_layers is not None:
return FullIterationSpace.create_with_ghost_layers( return FullIterationSpace.create_with_ghost_layers(
ctx, archetype_field, ctx.options.ghost_layers ctx, archetype_field, ghost_layers
) )
elif ctx.options.iteration_slice is not None: elif iteration_slice is not None:
raise PsInternalCompilerError("Iteration slices not supported yet") raise PsInternalCompilerError("Iteration slices not supported yet")
else: else:
return FullIterationSpace.create_with_ghost_layers( return FullIterationSpace.create_with_ghost_layers(
......
...@@ -158,7 +158,7 @@ class Typifier(Mapper): ...@@ -158,7 +158,7 @@ class Typifier(Mapper):
return var return var
def map_variable(self, var: pb.Variable, tc: TypeContext) -> PsTypedVariable: def map_variable(self, var: pb.Variable, tc: TypeContext) -> PsTypedVariable:
dtype = self._ctx.options.default_dtype dtype = self._ctx.default_dtype
typed_var = PsTypedVariable(var.name, dtype) typed_var = PsTypedVariable(var.name, dtype)
self._apply_target_type(typed_var, dtype, tc) self._apply_target_type(typed_var, dtype, tc)
return typed_var return typed_var
...@@ -175,7 +175,7 @@ class Typifier(Mapper): ...@@ -175,7 +175,7 @@ class Typifier(Mapper):
def map_array_access(self, access: PsArrayAccess, tc: TypeContext) -> PsArrayAccess: def map_array_access(self, access: PsArrayAccess, tc: TypeContext) -> PsArrayAccess:
self._apply_target_type(access, access.dtype, tc) self._apply_target_type(access, access.dtype, tc)
index = self.rec( index = self.rec(
access.index_tuple[0], TypeContext(self._ctx.options.index_dtype) access.index_tuple[0], TypeContext(self._ctx.index_dtype)
) )
return PsArrayAccess(access.base_ptr, index) return PsArrayAccess(access.base_ptr, index)
......
from typing import Sequence
from dataclasses import dataclass
from .enums import Target
from .field import Field, FieldType
from .backend.jit import JitBase
from .backend.exceptions import PsOptionsError
from .backend.types import PsIntegerType, PsNumericType, PsIeeeFloatType
from .backend.kernelcreation.defaults import Sympy as SpDefaults
from .backend.ast import PsKernelFunction from .backend.ast import PsKernelFunction
from .backend.kernelcreation import KernelCreationContext, KernelAnalysis, FreezeExpressions, Typifier from .backend.kernelcreation import (
KernelCreationContext,
KernelAnalysis,
FreezeExpressions,
Typifier,
)
from .backend.kernelcreation.iteration_space import ( from .backend.kernelcreation.iteration_space import (
create_sparse_iteration_space, create_sparse_iteration_space,
create_full_iteration_space, create_full_iteration_space,
...@@ -7,24 +24,126 @@ from .backend.kernelcreation.iteration_space import ( ...@@ -7,24 +24,126 @@ from .backend.kernelcreation.iteration_space import (
from .backend.kernelcreation.transformations import EraseAnonymousStructTypes from .backend.kernelcreation.transformations import EraseAnonymousStructTypes
from .enums import Target from .enums import Target
from .config import CreateKernelConfig
from .sympyextensions import AssignmentCollection from .sympyextensions import AssignmentCollection
@dataclass
class CreateKernelConfig:
"""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, ...
"""
jit: JitBase | None = None
"""Just-in-time compiler used to compile and load the kernel for invocation from the current Python environment.
If left at `None`, a default just-in-time compiler will be inferred from the `target` parameter.
To explicitly disable JIT compilation, pass `nbackend.jit.no_jit`.
"""
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_field: Field | None = None
"""Index field for a sparse kernel.
If this option is set, a sparse kernel with the given field as index field will be generated.
"""
"""Data Types"""
index_dtype: PsIntegerType = SpDefaults.index_dtype
"""Data type used for all index calculations."""
default_dtype: PsNumericType = PsIeeeFloatType(64)
"""Default numeric data type.
This data type will be applied to all untyped symbols.
"""
def __post_init__(self):
# Check iteration space argument consistency
if (
int(self.iteration_slice is not None)
+ int(self.ghost_layers is not None)
+ int(self.index_field is not None)
> 1
):
raise PsOptionsError(
"Parameters `iteration_slice`, `ghost_layers` and 'index_field` are mutually exclusive; "
"at most one of them may be set."
)
# Check index field
if (
self.index_field is not None
and self.index_field.field_type != FieldType.INDEXED
):
raise PsOptionsError(
"Only fields with `field_type == FieldType.INDEXED` can be specified as `index_field`"
)
# Infer JIT
if self.jit is None:
match self.target:
case Target.CPU:
from .backend.jit import legacy_cpu
self.jit = legacy_cpu
case _:
raise NotImplementedError(
f"No default JIT compiler implemented yet for target {self.target}"
)
def create_kernel( def create_kernel(
assignments: AssignmentCollection, assignments: AssignmentCollection,
config: CreateKernelConfig = CreateKernelConfig(), config: CreateKernelConfig = CreateKernelConfig(),
): ):
"""Create a kernel AST from an assignment collection.""" """Create a kernel AST from an assignment collection."""
ctx = KernelCreationContext(config) ctx = KernelCreationContext(
default_dtype=config.default_dtype, index_dtype=config.index_dtype
)
analysis = KernelAnalysis(ctx) analysis = KernelAnalysis(ctx)
analysis(assignments) analysis(assignments)
if len(ctx.fields.index_fields) > 0 or ctx.options.index_field is not None: if len(ctx.fields.index_fields) > 0 or config.index_field is not None:
ispace = create_sparse_iteration_space(ctx, assignments) ispace = create_sparse_iteration_space(
ctx, assignments, index_field=config.index_field
)
else: else:
ispace = create_full_iteration_space(ctx, assignments) ispace = create_full_iteration_space(
ctx,
assignments,
ghost_layers=config.ghost_layers,
iteration_slice=config.iteration_slice,
)
ctx.set_iteration_space(ispace) ctx.set_iteration_space(ispace)
...@@ -55,7 +174,9 @@ def create_kernel( ...@@ -55,7 +174,9 @@ def create_kernel(
kernel_ast = platform.optimize(kernel_ast) kernel_ast = platform.optimize(kernel_ast)
assert config.jit is not None assert config.jit is not None
function = PsKernelFunction(kernel_ast, config.target, name=config.function_name, jit=config.jit) function = PsKernelFunction(
kernel_ast, config.target, name=config.function_name, jit=config.jit
)
function.add_constraints(*ctx.constraints) function.add_constraints(*ctx.constraints)
return function return function
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment