For developers: AST Nodes and Transformations

AST Nodes

class Node(parent=None)

Base class for all AST nodes.

property args: List[Node | Expr]

Returns all arguments/children of this node.

property symbols_defined: Set[Symbol]

Set of symbols which are defined by this node.

property undefined_symbols: Set[Symbol]

Symbols which are used but are not defined inside this node.

subs(subs_dict)

Inplace! Substitute, similar to sympy’s but modifies the AST inplace.

Return type:

None

atoms(arg_type)

Returns a set of all descendants recursively, which are an instance of the given type.

Return type:

Set[Any]

class Conditional(condition_expr, true_block, false_block=None)

Conditional that maps to a ‘if’ statement in C/C++.

Try to avoid using this node inside of loops, since currently this construction can not be vectorized. Consider using assignments with sympy.Piecewise in this case.

Parameters:
  • condition_expr (Basic) – sympy relational expression

  • true_block (Union[Block, SympyAssignment]) – block which is run if conditional is true

  • false_block (Optional[Block]) – optional block which is run if conditional is false

subs(subs_dict)

Inplace! Substitute, similar to sympy’s but modifies the AST inplace.

property args

Returns all arguments/children of this node.

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

replace_by_true_block()

Replaces the conditional by its True block

replace_by_false_block()

Replaces the conditional by its False block

class KernelFunction(body, target, backend, compile_function, ghost_layers, function_name='kernel', assignments=None)
class Parameter(symbol, fields)

Function parameter.

Each undefined symbol in a KernelFunction node becomes a parameter to the function. Parameters are either symbols introduced by the user that never occur on the left hand side of an Assignment, or are related to fields/arrays passed to the function.

A parameter consists of the typed symbol (symbol property). For field related parameters this is a symbol defined in pystencils.kernelparameters. If the parameter is related to one or multiple fields, these fields are referenced in the fields property.

property target

See pystencils.Target

property backend

Backend

Type:

Backend for generating the code

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

property args

Returns all arguments/children of this node.

property fields_accessed: Set[Field]

fields which are accessed inside this kernel function

Type:

Set of Field instances

get_parameters()

Returns list of parameters for this function.

This function is expensive, cache the result where possible!

Return type:

Sequence[Parameter]

class SkipIteration(parent=None)
property args

Returns all arguments/children of this node.

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

class Block(nodes)
property args

Returns all arguments/children of this node.

subs(subs_dict)

Inplace! Substitute, similar to sympy’s but modifies the AST inplace.

Return type:

None

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

class PragmaBlock(pragma_line, nodes)
class LoopOverCoordinate(body, coordinate_to_loop_over, start, stop, step=1, is_block_loop=False, custom_loop_ctr=None)
subs(subs_dict)

Inplace! Substitute, similar to sympy’s but modifies the AST inplace.

property args

Returns all arguments/children of this node.

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

class SympyAssignment(lhs_symbol, rhs_expr, is_const=True, use_auto=False)
subs(subs_dict)

Inplace! Substitute, similar to sympy’s but modifies the AST inplace.

property args

Returns all arguments/children of this node.

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

class ResolvedFieldAccess(base, linearized_index, field, offsets, idx_coordinate_values)
class TemporaryMemoryAllocation(typed_symbol, size, align_offset)

Node for temporary memory buffer allocation.

Always allocates aligned memory.

Parameters:
  • typed_symbol (TypedSymbol) – symbol used as pointer (has to be typed)

  • size – number of elements to allocate

  • align_offset – the align_offset’s element is aligned

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

property args

Returns all arguments/children of this node.

offset(byte_alignment)

Number of ELEMENTS to skip for a pointer that is aligned to byte_alignment.

class TemporaryMemoryFree(alloc_node)
property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

property args

Returns all arguments/children of this node.

class SourceCodeComment(text)
property args

Returns all arguments/children of this node.

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

class EmptyLine
property args

Returns all arguments/children of this node.

property symbols_defined

Set of symbols which are defined by this node.

property undefined_symbols

Symbols which are used but are not defined inside this node.

class ConditionalFieldAccess(*args)

pystencils.Field.Access that is only executed if a certain condition is met. Can be used, for instance, for out-of-bound checks.

Transformations

class NestedScopes

Symbol visibility model using nested scopes

  • every accessed symbol that was not defined before, is added as a “free parameter”

  • free parameters are global, i.e. they are not in scopes

  • push/pop adds or removes a scope

>>> s = NestedScopes()
>>> s.access_symbol("a")
>>> s.is_defined("a")
False
>>> s.free_parameters
{'a'}
>>> s.define_symbol("b")
>>> s.is_defined("b")
True
>>> s.push()
>>> s.is_defined_locally("b")
False
>>> s.define_symbol("c")
>>> s.pop()
>>> s.is_defined("c")
False
iterate_loops_by_depth(node, nesting_depth)

Iterate all LoopOverCoordinate nodes in the given AST of the specified nesting depth.

Parameters:
  • node – Root node of the abstract syntax tree

  • nesting_depth – Nesting depth of the loops the pragmas should be applied to. Outermost loop has depth 0. A depth of -1 indicates the innermost loops.

Returns: Iterable listing all loop nodes of given nesting depth.

unify_shape_symbols(body, common_shape, fields)

Replaces symbols for array sizes to ensure they are represented by the same unique symbol.

When creating a kernel with variable array sizes, all passed arrays must have the same size. This is ensured when the kernel is called. Inside the kernel this means that only on symbol has to be used instead of one for each field. For example shape_arr1[0] and shape_arr2[0] must be equal, so they should also be represented by the same symbol.

Parameters:
  • body – ast node, for the kernel part where substitutions is made, is modified in-place

  • common_shape – shape of the field that was chosen

  • fields – all fields whose shapes should be replaced by common_shape

get_common_field(field_set)

Takes a set of pystencils Fields, checks if a common spatial shape exists and returns one representative field, that can be used for shape information etc. in the kernel creation. If the fields have different shapes ValueError is raised

make_loop_over_domain(body, iteration_slice=None, ghost_layers=None, loop_order=None)

Uses pystencils.field.Field.Access to create (multiple) loops around given AST.

Parameters:
  • body – Block object with inner loop contents

  • iteration_slice – if not None, iteration is done only over this slice of the field

  • ghost_layers – a sequence of pairs for each coordinate with lower and upper nr of ghost layers if None, the number of ghost layers is determined automatically and assumed to be equal for a all dimensions

  • loop_order – loop ordering from outer to inner loop (optimal ordering is same as layout)

Returns:

tuple of loop-node, ghost_layer_info

create_intermediate_base_pointer(field_access, coordinates, previous_ptr)

Addressing elements in structured arrays is done with \(ptr\left[ \sum_i c_i \cdot s_i \right]\) where \(c_i\) is the coordinate value and \(s_i\) the stride of a coordinate. The sum can be split up into multiple parts, such that parts of it can be pulled before loops. This function creates such an access for coordinates \(i \in \mbox{coordinates}\). Returns a new typed symbol, where the name encodes which coordinates have been resolved.

Parameters:
  • field_access – instance of pystencils.field.Field.Access which provides strides and offsets

  • coordinates – mapping of coordinate ids to its value, where stride*value is calculated

  • previous_ptr – the pointer which is de-referenced

Returns

tuple with the new pointer symbol and the calculated offset

Examples

>>> field = Field.create_generic('myfield', spatial_dimensions=2, index_dimensions=1)
>>> x, y = sp.symbols("x y")
>>> prev_pointer = TypedSymbol("ptr", "double")
>>> create_intermediate_base_pointer(field[1,-2](5), {0: x}, prev_pointer)
(ptr_01, _stride_myfield_0*x + _stride_myfield_0)
>>> create_intermediate_base_pointer(field[1,-2](5), {0: x, 1 : y }, prev_pointer)
(ptr_01_1m2, _stride_myfield_0*x + _stride_myfield_0 + _stride_myfield_1*y - 2*_stride_myfield_1)
parse_base_pointer_info(base_pointer_specification, loop_order, spatial_dimensions, index_dimensions)

Creates base pointer specification for resolve_field_accesses() function.

Specification of how many and which intermediate pointers are created for a field access. For example [ (0), (2,3,)] creates on base pointer for coordinates 2 and 3 and writes the offset for coordinate zero directly in the field access. These specifications are defined dependent on the loop ordering. This function translates more readable version into the specification above.

Allowed specifications:
  • “spatialInner<int>” spatialInner0 is the innermost loop coordinate, spatialInner1 the loop enclosing the innermost

  • “spatialOuter<int>” spatialOuter0 is the outermost loop

  • “index<int>”: index coordinate

  • “<int>”: specifying directly the coordinate

Parameters:
  • base_pointer_specification – nested list with above specifications

  • loop_order – list with ordering of loops from outer to inner

  • spatial_dimensions – number of spatial dimensions

  • index_dimensions – number of index dimensions

Returns:

list of tuples that can be passed to resolve_field_accesses()

Examples

>>> parse_base_pointer_info([['spatialOuter0'], ['index0']], loop_order=[2,1,0],
...                         spatial_dimensions=3, index_dimensions=1)
[[0], [3], [1, 2]]
get_base_buffer_index(ast_node, loop_counters=None, loop_iterations=None)

Used for buffer fields to determine the linearized index of the buffer dependent on loop counter symbols.

Parameters:
  • ast_node – ast before any field accesses are resolved

  • loop_counters – for CPU kernels: leave to default ‘None’ (can be determined from loop nodes) for GPU kernels: list of ‘loop counters’ from inner to outer loop

  • loop_iterations – iteration slice for each loop from inner to outer, for CPU kernels leave to default

Returns:

base buffer index - required by ‘resolve_buffer_accesses’ function

resolve_field_accesses(ast_node, read_only_field_names=None, field_to_base_pointer_info=mappingproxy({}), field_to_fixed_coordinates=mappingproxy({}))

Substitutes pystencils.field.Field.Access nodes by array indexing

Parameters:
  • ast_node – the AST root

  • read_only_field_names – set of field names which are considered read-only

  • field_to_base_pointer_info – a list of tuples indicating which intermediate base pointers should be created for details see parse_base_pointer_info()

  • field_to_fixed_coordinates – map of field name to a tuple of coordinate symbols. Instead of using the loop counters to index the field these symbols are used as coordinates

Returns

transformed AST

move_constants_before_loop(ast_node)

Moves pystencils.ast.SympyAssignment nodes out of loop body if they are iteration independent.

Call this after creating the loop structure with make_loop_over_domain()

split_inner_loop(ast_node, symbol_groups)

Splits inner loop into multiple loops to minimize the amount of simultaneous load/store streams

Parameters:
  • ast_node (Node) – AST root

  • symbol_groups – sequence of symbol sequences: for each symbol sequence a new inner loop is created which updates these symbols and their dependent symbols. Symbols which are in none of the symbolGroups and which no symbol in a symbol group depends on, are not updated!

cut_loop(loop_node, cutting_points)

Cuts loop at given cutting points.

One loop is transformed into len(cuttingPoints)+1 new loops that range from old_begin to cutting_points[1], …, cutting_points[-1] to old_end

Modifies the ast in place. Note Issue #5783 of SymPy. Deepcopy will evaluate mul https://github.com/sympy/sympy/issues/5783

Returns:

list of new loop nodes

simplify_conditionals(node, loop_counter_simplification=False)

Removes conditionals that are always true/false.

Parameters:
  • node (Node) – ast node, all descendants of this node are simplified

  • loop_counter_simplification (bool) – if enabled, tries to detect if a conditional is always true/false depending on the surrounding loop. For example if the surrounding loop goes from x=0 to 10 and the condition is x < 0, it is removed. This analysis needs the integer set library (ISL) islpy, so it is not done by default.

Return type:

None

cleanup_blocks(node)

Curly Brace Removal: Removes empty blocks, and replaces blocks with a single child by its child

Return type:

None

remove_conditionals_in_staggered_kernel(function_node, include_first=True)

Removes conditionals of a kernel that iterates over staggered positions by splitting the loops at last or first and last element

Return type:

None

get_optimal_loop_ordering(fields)

Determines the optimal loop order for a given set of fields. If the fields have different memory layout or different sizes an exception is thrown.

Parameters:

fields – sequence of fields

Returns:

list of coordinate ids, where the first list entry should be the outermost loop

get_loop_hierarchy(ast_node)

Determines the loop structure around a given AST node, i.e. the node has to be inside the loops.

Returns:

sequence of LoopOverCoordinate nodes, starting from outer loop to innermost loop

get_loop_counter_symbol_hierarchy(ast_node)

Determines the loop counter symbols around a given AST node. :type ast_node: :param ast_node: the AST node :return: list of loop counter symbols, where the first list entry is the symbol of the innermost loop

replace_inner_stride_with_one(ast_node)

Replaces the stride of the innermost loop of a variable sized kernel with 1 (assumes optimal loop ordering).

Variable sized kernels can handle arbitrary field sizes and field shapes. However, the kernel is most efficient if the innermost loop accesses the fields with stride 1. The inner loop can also only be vectorized if the inner stride is 1. This transformation hard codes this inner stride to one to enable e.g. vectorization.

Warning: the assumption is not checked at runtime!

Return type:

None

loop_blocking(ast_node, block_size)

Blocking of loops to enhance cache locality. Modifies the ast node in-place.

Parameters:
  • ast_node (KernelFunction) – kernel function node before vectorization transformation has been applied

  • block_size – sequence defining block size in x, y, (z) direction. If chosen as zero the direction will not be used for blocking.

Return type:

int

Returns:

number of dimensions blocked