AST Transformations

pystencils.backend.transformations

This module contains various transformation and optimization passes that can be executed on the backend AST.

Canonical Form

Many transformations in this module require that their input AST is in canonical form. This means that:

  • Each symbol, constant, and expression node is annotated with a data type;

  • Each symbol has at most one declaration;

  • Each symbol that is never written to apart from its declaration has a const type; and

  • Each symbol whose type is not const has at least one non-declaring assignment.

The first requirement can be ensured by running the Typifier on each newly constructed subtree. The other three requirements are ensured by the CanonicalizeSymbols pass, which should be run first before applying any optimizing transformations. All transformations in this module retain canonicality of the AST.

Canonicality allows transformations to forego various checks that would otherwise be necessary to prove their legality.

Certain transformations, like the auto-vectorizer (TODO), state additional requirements, e.g. the absence of loop-carried dependencies.

Transformations

Canonicalization

class pystencils.backend.transformations.CanonicalizeSymbols(ctx, constify=True)

Remove duplicate symbol declarations and declare all non-updated symbols const.

The CanonicalizeSymbols pass will remove multiple declarations of the same symbol by renaming all but the last occurence, and will optionally const-qualify all symbols encountered in the AST that are never updated.

Parameters:
__call__(node)

Call self as a function.

Return type:

PsAstNode

Parameters:

node (PsAstNode) –

AST Cloning

class pystencils.backend.transformations.CanonicalClone(ctx)

Clone a subtree, and rename all symbols declared inside it to retain canonicality.

Parameters:

ctx (KernelCreationContext) –

__call__(node)

Call self as a function.

Return type:

TypeVar(Node_T, bound= PsAstNode)

Parameters:

node (Node_T) –

Simplifying Transformations

class pystencils.backend.transformations.EliminateConstants(ctx, extract_constant_exprs=False)

Eliminate constant expressions in various ways.

  • Constant folding: Nontrivial constant integer (and optionally floating point) expressions are evaluated and replaced by their result

  • Idempotence elimination: Idempotent operations (e.g. addition of zero, multiplication with one) are replaced by their result

  • Dominance elimination: Multiplication by zero is replaced by zero

  • Constant extraction: Optionally, nontrivial constant expressions are extracted and listed at the beginning of the outermost block.

Parameters:
__call__(node)

Call self as a function.

Return type:

PsAstNode

Parameters:

node (PsAstNode) –

class pystencils.backend.transformations.EliminateBranches(ctx)

Replace conditional branches by their then- or else-branch if their condition can be unequivocally evaluated.

This pass will attempt to evaluate branch conditions within their context in the AST, and replace conditionals by either their then- or their else-block if the branch is unequivocal.

TODO: If islpy is installed, this pass will incorporate information about the iteration regions of enclosing loops into its analysis.

Parameters:

ctx (KernelCreationContext) –

__call__(node)

Call self as a function.

Return type:

PsAstNode

Parameters:

node (PsAstNode) –

Code Motion

class pystencils.backend.transformations.HoistLoopInvariantDeclarations(ctx)

Hoist loop-invariant declarations out of the loop nest.

This transformation moves loop-invariant symbol declarations outside of the loop nest to prevent their repeated execution within the loops. If this transformation results in the complete elimination of a loop body, the respective loop is removed.

HoistLoopInvariantDeclarations assumes that symbols are canonical; in particular, each symbol may have at most one declaration. To ensure this, a CanonicalizeSymbols pass should be run before HoistLoopInvariantDeclarations.

HoistLoopInvariantDeclarations assumes that all PsMathFunction s are pure (have no side effects), but makes no such assumption about instances of CFunction.

Parameters:

ctx (KernelCreationContext) –

__call__(node)

Call self as a function.

Return type:

PsAstNode

Parameters:

node (PsAstNode) –

Loop Reshaping Transformations

class pystencils.backend.transformations.ReshapeLoops(ctx)

Various transformations for reshaping loop nests.

Parameters:

ctx (KernelCreationContext) –

peel_loop_front(loop, num_iterations, omit_range_check=False)

Peel off iterations from the front of a loop.

Removes num_iterations from the front of the given loop and returns them as a sequence of independent blocks.

Parameters:

  • loop (PsLoop) – The loop node from which to peel iterations

  • num_iterations (int) – The number of iterations to peel off

  • omit_range_check (bool) – If set to True, assume that the peeled-off iterations will always be executed, and omit their enclosing conditional.

Return type:

tuple[Sequence[PsBlock], PsLoop]

Returns:

Tuple containing the peeled-off iterations as a sequence of blocks, and the remaining loop.

cut_loop(loop, cutting_points)

Cut a loop at the given cutting points.

Cut the given loop at the iterations specified by the given cutting points, producing n new subtrees representing the iterations (loop.start:cutting_points[0]), (cutting_points[0]:cutting_points[1]), ..., (cutting_points[-1]:loop.stop).

Resulting subtrees representing zero iterations are dropped; subtrees representing exactly one iteration are returned without the trivial loop structure.

Currently, cut_loop performs no checks to ensure that the given cutting points are in fact inside the loop’s iteration range.

Return type:

Sequence[PsLoop | PsBlock]

Returns:

Sequence of n subtrees representing the respective iteration ranges

Parameters:

Code Lowering and Materialization

class pystencils.backend.transformations.EraseAnonymousStructTypes(ctx)

Lower anonymous struct arrays to a byte-array representation.

For arrays whose element type is an anonymous struct, the struct type is erased from the base pointer, making it a pointer to uint8_t. Member lookups on accesses into these arrays are then transformed using type casts.

Parameters:

ctx (KernelCreationContext) –

__call__(node)

Call self as a function.

Return type:

PsAstNode

Parameters:

node (PsAstNode) –

class pystencils.backend.transformations.SelectFunctions(platform)

Traverse the AST to replace all instances of PsMathFunction by their implementation provided by the given Platform.

Parameters:

platform (Platform) –

__call__(node)

Call self as a function.

Return type:

PsAstNode

Parameters:

node (PsAstNode) –