Kernel Translation

Kernel Translation#

The Kernel Creation Context#

class pystencils.backend.kernelcreation.KernelCreationContext(default_dtype=PsIeeeFloatType(width=64, const=False), index_dtype=PsIntegerType(width=64, signed=True, const=False))#

Manages the translation process from the SymPy frontend to the backend AST, and collects all necessary information for the translation:

  • Data Types: The kernel creation context manages the default data types for loop limits and counters, index calculations, and the typifier.

  • Symbols: The context maintains a symbol table, keeping track of all symbols encountered during kernel translation together with their types.

  • Fields and Arrays: The context collects all fields encountered during code generation, applies a few consistency checks to them, and manages their associated arrays.

  • Iteration Space: The context manages the iteration space of the kernel currently being translated.

  • Constraints: The context collects all kernel parameter constraints introduced during the translation process.

  • Required Headers: The context collects all header files required for the kernel to run.

Parameters:
property default_dtype: PsNumericType#

Data type used by default for numerical expressions

property index_dtype: PsIntegerType#

Data type used by default for index expressions

resolve_dynamic_type(dtype)#

Selects the appropriate data type for DynamicType instances, and returns all other types as they are.

Return type:

PsType

Parameters:

dtype (DynamicType | PsType)

get_symbol(name, dtype=None)#

Retrieve the symbol with the given name and data type from the symbol table.

If no symbol named name exists, a new symbol with the given data type is created.

If a symbol with the given name already exists and dtype is not None, the given data type will be applied to it, and it is returned. If the symbol already has a different data type, an error will be raised.

If the symbol already exists and dtype is None, the existing symbol is returned without checking or altering its data type.

Parameters:
Return type:

PsSymbol

get_new_symbol(name, dtype=None)#

Always create a new symbol, deduplicating its name if another symbol with the same name already exists.

Return type:

PsSymbol

Parameters:
find_symbol(name)#

Find a symbol with the given name in the symbol table, if it exists.

Return type:

PsSymbol | None

Returns:

The symbol with the given name, or None if no such symbol exists.

Parameters:

name (str)

add_symbol(symbol)#

Add an existing symbol to the symbol table.

If a symbol with the same name already exists, an error will be raised.

Parameters:

symbol (PsSymbol)

replace_symbol(old, new)#

Replace one symbol by another.

The two symbols old and new must have the same name, but may have different data types.

Parameters:
duplicate_symbol(symb, new_dtype=None)#

Canonically duplicates the given symbol.

A new symbol with the new name symb.name + "__<counter>" and optionally a different data type is created, added to the symbol table, and returned. The counter reflects the number of previously created duplicates of this symbol.

Return type:

PsSymbol

Parameters:
basename(symb)#

Returns the original name a symbol had before duplication.

Return type:

str

Parameters:

symb (PsSymbol)

property symbols: Iterable[PsSymbol]#

Return an iterable of all symbols listed in the symbol table.

property fields: FieldsInKernel#

Collection of fields that occured during the current kernel translation.

add_field(field)#

Add the given field to the context’s fields collection.

This method adds the passed field to the context’s field collection, which is accesible through the field member, and creates the underlying buffer for the field which is retrievable through get_buffer. Before adding the field to the collection, various sanity and constraint checks are applied.

Parameters:

field (Field)

get_buffer(field)#

Retrieve the underlying array for a given field.

If the given field was not previously registered using add_field, this method internally calls add_field to check the field for consistency.

Return type:

PsBuffer

Parameters:

field (Field)

set_iteration_space(ispace)#

Set the iteration space used for the current kernel.

Parameters:

ispace (IterationSpace)

Analysis and Constraints Checks#

class pystencils.backend.kernelcreation.KernelAnalysis(ctx, check_access_independence=True, check_double_writes=True)#

General analysis pass over a kernel expressed using the SymPy frontend.

The kernel analysis fulfills two tasks. It checks the SymPy input for consistency, and populates the context with required knowledge about the kernel.

A KernelAnalysis object may be called at most once.

Consistency and Constraints

The following checks are performed:

  • SSA Form: The given assignments must be in single-assignment form; each symbol must be written at most once.

  • Independence of Accesses: To avoid loop-carried dependencies, each field may be written at most once at each index, and if a field is written at some location with index i, it may only be read with index i in the same location.

  • Independence of Writes: A weaker requirement than access independence; each field may only be written once at each index.

  • Dimension of index fields: Index fields occuring in the kernel must have exactly one spatial dimension.

Knowledge Collection

The following knowledge is collected into the context:

  • The set of fields accessed in the kernel

Parameters:
class FieldAndIndex(field, index)#
field#

Alias for field number 0

index#

Alias for field number 1

SymPy Parsing and IR Construction#

class pystencils.backend.kernelcreation.AstFactory(ctx)#

Factory providing a convenient interface for building syntax trees.

The AstFactory uses the defaults provided by the given KernelCreationContext to quickly create AST nodes. Depending on context (numerical, loop indexing, etc.), symbols and constants receive either ctx.default_dtype or ctx.index_dtype.

Parameters:

ctx (KernelCreationContext) – The kernel creation context

parse_sympy(sp_obj)#

Parse a SymPy expression or assignment through FreezeExpressions and Typifier.

The expression or assignment will be typified in a numerical context, using the kernel creation context’s default_dtype.

Parameters:

sp_obj (Expr | Tuple | Relational | BooleanFunction | AssignmentBase) – A SymPy expression or assignment

Return type:

PsAstNode

parse_index(idx)#

Parse the given object as an expression with data type ctx.index_dtype.

Parameters:

idx (PsExpression | PsSymbol | PsConstant | Expr | int | integer)

parse_slice(iter_slice, normalize_to=None)#

Parse a slice to obtain start, stop and step expressions for a loop or iteration space dimension.

The slice entries may be instances of PsExpression, PsSymbol or PsConstant, in which case they must typify with the kernel creation context’s index_dtype. They may also be sympy expressions or integer constants, in which case they are parsed to AST objects and must also typify with the kernel creation context’s index_dtype.

The step member of the slice, if it is constant, must be positive.

The slice may optionally be normalized with respect to an upper iteration limit. If normalize_to is specified, negative integers in iter_slice.start and iter_slice.stop will be added to that normalization limit.

Parameters:
Return type:

tuple[PsExpression, PsExpression, PsExpression]

loop(ctr_name, iteration_slice, body)#

Create a loop from a slice.

Parameters:

  • ctr_name (str) – Name of the loop counter

  • iteration_slice (slice) – The iteration region as a slice; see parse_slice.

  • body (PsBlock) – The loop body

loop_nest(counters, slices, body)#

Create a loop nest from a sequence of slices.

Example: This snippet creates a 3D loop nest with ten iterations in each dimension:

>>> from pystencils import make_slice
>>> ctx = KernelCreationContext()
>>> factory = AstFactory(ctx)
>>> loop = factory.loop_nest(("i", "j", "k"), make_slice[:10,:10,:10], PsBlock([]))
Parameters:
Return type:

PsLoop

loops_from_ispace(ispace, body, loop_order=None)#

Create a loop nest from a dense iteration space.

Parameters:
Return type:

PsLoop

class pystencils.backend.kernelcreation.FreezeExpressions(ctx)#

Convert expressions and kernels expressed in the SymPy language to the code generator’s internal representation.

This class accepts a subset of the SymPy symbolic algebra language complemented with the extensions implemented in pystencils.sympyextensions, and converts it to the abstract syntax tree representation of the pystencils code generator. It is invoked early during the code generation process.

TODO: Document the full set of supported SymPy features, with restrictions and caveats TODO: Properly document the SymPy extensions provided by pystencils

Parameters:

ctx (KernelCreationContext)

map_Function(func)#

Map SymPy function calls by mapping sympy function classes to backend-supported function symbols.

If applicable, functions are mapped to binary operators, e.g. PsBitwiseXor. Other SymPy functions are frozen to an instance of PsFunction.

Return type:

PsExpression

Parameters:

func (Function)

Type Checking and Inference#

class pystencils.backend.kernelcreation.typification.TypeContext(target_type=None)#

Typing context, with support for type inference and checking.

Instances of this class are used to propagate and check data types across expression subtrees of the AST. Each type context has a target type target_type, which shall be applied to all expressions it covers.

Just like the types of expressions, the context’s target type must never be const. This is ensured by this class, which removes const-qualification from the target type when it is set.

Parameters:

target_type (PsType | None)

property target_type: PsType | None#

The target type of this type context.

add_hook(hook)#

Adds a resolution hook to this context.

The hook will be called with the context’s target type as soon as it becomes known, which might be immediately.

Parameters:

hook (Callable[[PsType], None])

apply_dtype(dtype, expr=None)#

Applies the given dtype to this type context, and optionally to the given expression.

If the context’s target_type is already known, it must be compatible with the given dtype. If the target type is still unknown, target_type is set to dtype and retroactively applied to all deferred expressions.

If an expression is specified, it will be covered by the type context. If the expression already has a data type set, it must be compatible with the target type and will be replaced by it.

Parameters:
infer_dtype(expr)#

Infer the data type for the given expression.

If the target_type of this context is already known, it will be applied to the given expression. Otherwise, the expression is deferred, and a type will be applied to it as soon as apply_dtype is called on this context.

If the expression already has a data type set, it must be compatible with the target type and will be replaced by it.

Parameters:

expr (PsExpression)

class pystencils.backend.kernelcreation.Typifier(ctx)#

Apply data types to expressions.

Contextual Typing

The Typifier will traverse the AST and apply a contextual typing scheme to figure out the data types of all encountered expressions. To this end, it covers each expression tree with a set of disjoint typing contexts. All nodes covered by the same typing context must have the same type.

Starting from an expression’s root, a typing context is implicitly expanded through the recursive descent into a node’s children. In particular, a child is typified within the same context as its parent if the node’s semantics require parent and child to have the same type (e.g. at arithmetic operators, mathematical functions, etc.). If a node’s child is required to have a different type, a new context is opened.

For each typing context, its target type is prescribed by the first node encountered during traversal whose type is fixed according to its typing rules. All other nodes covered by the context must share that type.

The types of arithmetic operators, mathematical functions, and untyped constants are inferred from their context’s target type. If one of these is encountered while no target type is set yet in the context, the expression is deferred by storing it in the context, and will be assigned a type as soon as the target type is fixed.

Typing Rules

The following general rules apply:

  • The context’s default_dtype is applied to all untyped symbols encountered inside a right-hand side expression

  • If an untyped symbol is encountered on an assignment’s left-hand side, it will first be attempted to infer its type from the right-hand side. If that fails, the context’s default_dtype will be applied.

  • It is an error if an untyped symbol occurs in the same type context as a typed symbol or constant with a non-default data type.

  • By default, all expressions receive a const type unless they occur on a (non-declaration) assignment’s left-hand side

Typing of symbol expressions

Some expressions (PsSymbolExpr, PsBufferAcc) encapsulate symbols and inherit their data types.

Parameters:

ctx (KernelCreationContext)

visit(node)#

Recursive processing of structural nodes

Return type:

None

Parameters:

node (PsAstNode)

visit_expr(expr, tc)#

Recursive processing of expression nodes.

This method opens, expands, and closes typing contexts according to the respective expression’s typing rules. It may add or check restrictions only when opening or closing a type context.

The actual type inference and checking during context expansion are performed by the methods of TypeContext. visit_expr tells the typing context how to handle an expression by calling either apply_dtype or infer_dtype.

Return type:

None

Parameters:
Contents