Constants, Memory Objects, and Functions

Constants, Memory Objects, and Functions#

Memory Objects: Symbols and Buffers#

The Memory Model#

In order to reason about memory accesses, mutability, invariance, and aliasing, the pystencils backend uses a very simple memory model. There are three types of memory objects:

  • Symbols (PsSymbol), which act as registers for data storage within the scope of a kernel

  • Field buffers (PsBuffer), which represent a contiguous block of memory the kernel has access to, and

  • the unmanaged heap, which is a global catch-all memory object which all pointers not belonging to a field array point into.

All of these objects are disjoint, and cannot alias each other. Each symbol exists in isolation, field buffers do not overlap, and raw pointers are assumed not to point into memory owned by a symbol or field array. Instead, all raw pointers point into unmanaged heap memory, and are assumed to always alias one another: Each change brought to unmanaged memory by one raw pointer is assumed to affect the memory pointed to by another raw pointer.

Symbols#

In the pystencils IR, instances of PsSymbol represent what is generally known as “virtual registers”. These are memory locations that are private to a function, cannot be aliased or pointed to, and will finally reside either in physical registers or on the stack. Each symbol has a name and a data type. The data type may initially be None, in which case it should soon after be determined by the Typifier.

Other than their front-end counterpart sympy.Symbol, PsSymbol instances are mutable; their properties can and often will change over time. As a consequence, they are not comparable by value: two PsSymbol instances with the same name and data type will in general not be equal. In fact, most of the time, it is an error to have two identical symbol instances active.

Creating Symbols#

During kernel translation, symbols never exist in isolation, but should always be managed by a KernelCreationContext. Symbols can be created and retrieved using add_symbol and find_symbol. A symbol can also be duplicated using duplicate_symbol, which assigns a new name to the symbol’s copy. The KernelCreationContext keeps track of all existing symbols during a kernel translation run and makes sure that no name and data type conflicts may arise.

Never call the constructor of PsSymbol directly unless you really know what you are doing.

Symbol Properties#

Symbols can be annotated with arbitrary information using symbol properties. Each symbol property type must be a subclass of PsSymbolProperty. It is strongly recommended to implement property types using frozen dataclasses. For example, this snippet defines a property type that models pointer alignment requirements:

@dataclass(frozen=True)
class AlignmentProperty(UniqueSymbolProperty)
    """Require this pointer symbol to be aligned at a particular byte boundary."""
    
    byte_boundary: int

Inheriting from UniqueSymbolProperty ensures that at most one property of this type can be attached to a symbol at any time. Properties can be added, queried, and removed using the PsSymbol properties API listed below.

Many symbol properties are more relevant to consumers of generated kernels than to the code generator itself. The above alignment property, for instance, may be added to a pointer symbol by a vectorization pass to document its assumption that the pointer be properly aligned, in order to emit aligned load and store instructions. It then becomes the responsibility of the runtime system embedding the kernel to check this prequesite before calling the kernel. To make sure this information becomes visible, any properties attached to symbols exposed as kernel parameters will also be added to their respective Parameter instance.

Buffers#

Buffers, as represented by the PsBuffer class, represent contiguous, n-dimensional, linearized cuboid blocks of memory. Each buffer has a fixed name and element data type, and will be represented in the IR via three sets of symbols:

  • The base pointer is a symbol of pointer type which points into the buffer’s underlying memory area. Each buffer has at least one, its primary base pointer, whose pointed-to type must be the same as the buffer’s element type. There may be additional base pointers pointing into subsections of that memory. These additional base pointers may also have deviating data types, as is for instance required for type erasure in certain cases. To communicate its role to the code generation system, each base pointer needs to be marked as such using the BufferBasePtr property, .

  • The buffer shape defines the size of the buffer in each dimension. Each shape entry is either a symbol <PsSymbol> or a constant.

  • The buffer strides define the step size to go from one entry to the next in each dimension. Like the shape, each stride entry is also either a symbol or a constant.

The shape and stride symbols must all have the same data type, which will be stored as the buffer’s index data type.

Creating and Managing Buffers#

Similarily to symbols, buffers are typically managed by the KernelCreationContext, which associates each buffer to a front-end Field. Buffers for fields can be obtained using get_buffer. The context makes sure to avoid name conflicts between buffers.

Constants#

In the pystencils IR, numerical constants are represented by the PsConstant class. It interacts with the type system (in particular, with PsNumericType and its subclasses) to facilitate bit-exact storage, arithmetic, and type conversion of constants.

Each constant has a value and a data type. As long as the data type is None, the constant is untyped and its value may be any Python object. To add a data type, an instance of PsNumericType must either be set in the constructor, or be applied by converting an existing constant using interpret_as. Once a data type is set, the set of legal values is constrained by that type.

To facilitate the correctness of the internal representation, PsConstant calls PsNumericType.create_constant. This method must be overridden by subclasses of PsNumericType; it either returns an object that represents the numerical constant according to the rules of the data type, or raises an exception if that is not possible. The fixed-width integers, the IEEE-754 floating point types, and the corresponding vector variants that are implemented in pystencils use NumPy for this purpose.

The same protocol is used for type conversion of constants, using PsConstant.reinterpret_as.

Functions#

The pystencils IR models two primary kinds of functions: IR functions and external functions.

IR functions are functions that are only serve as an intermediate representation of a function, and must be lowered to a concrete implementation at some point during kernel translation. This is typically done by the SelectFunctions pass, in combination with the active Platform class.

IR function classes derive from the common base class PsIrFunction. The following groups of IR functions exist:

External functions with a fixed C-like signature are modelled by the CFunction class. They are used to inject platform-specific runtime APIs, vector intrinsics, and user-defined external functions into the IR. These are the only functions allowed to remain in a kernel by the time it is exported as C code.

Side Effects#

PsMathFunction and PsConstantFunction represent pure functions. Their occurences may be moved, optimized, or eliminated by the code generator. For CFunction, on the other hand, side effects are conservatively assumed, such that these cannot be freely manipulated.

Literals#

In the pystencils IR, a literal is an expression string, with an associated data type, that is taken literally and printed out verbatim by the code generator. They are represented by the PsLiteral class, and are used to represent compiler-builtins (like the CUDA variables threadIdx, blockIdx, …), preprocessor macros (like INFINITY), and other pieces of code that could not otherwise be modelled. Literals are assumed to be constant with respect to the kernel, and their evaluation is assumed to be free of side effects.

API Documentation#

class pystencils.codegen.properties.PsSymbolProperty#

Base class for symbol properties, which can be used to add additional information to symbols

class pystencils.codegen.properties.UniqueSymbolProperty#

Base class for unique properties, of which only one instance may be registered at a time.

class pystencils.codegen.properties.FieldShape(field, coordinate)#

Symbol acts as a shape parameter to a field.

Parameters:
class pystencils.codegen.properties.FieldStride(field, coordinate)#

Symbol acts as a stride parameter to a field.

Parameters:
class pystencils.codegen.properties.FieldBasePtr(field)#

Symbol acts as a base pointer to a field.

Parameters:

field (Field)

class pystencils.backend.memory.PsSymbol(name, dtype=None)#

A mutable symbol with name and data type.

Do not create objects of this class directly unless you know what you are doing; instead obtain them from a KernelCreationContext through KernelCreationContext.get_symbol. This way, the context can keep track of all symbols used in the translation run, and uniqueness of symbols is ensured.

Parameters:
apply_dtype(dtype)#

Apply the given data type to this symbol, raising a TypeError if it conflicts with a previously set data type.

Parameters:

dtype (PsType)

property properties: frozenset[PsSymbolProperty]#

Set of properties attached to this symbol

get_properties(prop_type)#

Retrieve all properties of the given type attached to this symbol

Return type:

set[PsSymbolProperty]

Parameters:

prop_type (type[PsSymbolProperty])

add_property(property)#

Attach a property to this symbol

Parameters:

property (PsSymbolProperty)

remove_property(property)#

Remove a property from this symbol. Does nothing if the property is not attached.

Parameters:

property (PsSymbolProperty)

class pystencils.backend.memory.BackendPrivateProperty#

Mix-in marker for symbol properties that are private to the backend and should not be exported to parameters

class pystencils.backend.memory.BufferBasePtr(buffer)#

Symbol acts as a base pointer to a buffer.

Parameters:

buffer (PsBuffer)

class pystencils.backend.memory.PsBuffer(name, element_type, base_ptr, shape, strides)#

N-dimensional contiguous linearized buffer in heap memory.

PsBuffer models the memory buffers underlying the Field class to the backend. Each buffer represents a contiguous block of memory that is non-aliased and disjoint from all other buffers.

Buffer shape and stride information are given either as constants or as symbols. All indexing expressions must have the same data type, which will be selected as the buffer’s index_dtype <PsBuffer.index_dtype>.

Each buffer has at least one base pointer, which can be retrieved via the PsBuffer.base_pointer property.

Parameters:
property name#

The buffer’s name

property base_pointer: PsSymbol#

Primary base pointer

property shape: list[PsSymbol | PsConstant]#

Buffer shape symbols and/or constants

property strides: list[PsSymbol | PsConstant]#

Buffer stride symbols and/or constants

property dim: int#

Dimensionality of this buffer

property index_type: PsIntegerType#

Index data type of this buffer; i.e. data type of its shape and stride symbols

property element_type: PsType#

Element type of this buffer

class pystencils.backend.constants.PsConstant(value, dtype=None)#

Type-safe representation of typed numerical constants.

This class models constants in the backend representation of kernels. A constant may be untyped, in which case its value may be any Python object.

If the constant is typed (i.e. its dtype is not None), its data type is used to check the validity of its value and to convert it into the type’s internal representation.

Instances of PsConstant are immutable.

Parameters:

  • value (Any) – The constant’s value

  • dtype (Optional[PsNumericType]) – The constant’s data type, or None if untyped.

interpret_as(dtype)#

Interprets this untyped constant with the given data type.

If this constant is already typed, raises an error.

Return type:

PsConstant

Parameters:

dtype (PsNumericType)

reinterpret_as(dtype)#

Reinterprets this constant with the given data type.

Other than interpret_as, this method also works on typed constants.

Return type:

PsConstant

Parameters:

dtype (PsNumericType)

property dtype: PsNumericType | None#

This constant’s data type, or None if it is untyped.

The data type of a constant always has const == True.

get_dtype()#

Retrieve this constant’s data type, throwing an exception if the constant is untyped.

Return type:

PsNumericType

class pystencils.backend.literals.PsLiteral(text, dtype)#

Representation of literal code.

Instances of this class represent code literals inside the AST. These literals are not to be confused with C literals; the name ‘Literal’ refers to the fact that the code generator takes them “literally”, printing them as they are.

Each literal has to be annotated with a type, and is considered constant within the scope of a kernel. Instances of PsLiteral are immutable.

The code generator assumes literals to be constant and pure with respect to the kernel: their evaluation at kernel runtime must not include any side effects.

Parameters:
class pystencils.backend.functions.PsFunction(name, num_args)#

Base class for functions occuring in the IR

Parameters:
property name: str#

Name of this function.

property arg_count: int#

Number of arguments this function takes

class pystencils.backend.functions.PsIrFunction(name, num_args)#

Base class for IR functions that must be lowered to target-specific implementations.

Parameters:
class pystencils.backend.functions.MathFunctions(value)#

Mathematical functions supported by the backend.

Each platform has to materialize these functions to a concrete implementation.

class pystencils.backend.functions.PsMathFunction(func)#

Homogeneously typed mathematical functions.

Parameters:

func (MathFunctions)

class pystencils.backend.functions.PsReductionWriteBack(reduction_op)#

Function representing a reduction kernel’s write-back step supported by the backend.

Each platform has to materialize this function to a concrete implementation.

Parameters:

reduction_op (ReductionOp)

class pystencils.backend.functions.ConstantFunctions(value)#

Numerical constant functions.

Each platform has to materialize these functions to a concrete implementation.

class pystencils.backend.functions.PsConstantFunction(func, dtype=None)#

Data-type-specific numerical constants.

Represents numerical constants which need to be exactly represented, e.g. transcendental numbers and non-finite constants.

Functions of this class are treated the same as PsConstant instances by most transforms. In particular, they are subject to the same contextual typing rules, and will be broadcast by the vectorizer.

Parameters:
property dtype: PsNumericType | None#

This constant function’s data type, or None if it is untyped.

get_dtype()#

Retrieve this constant function’s data type, throwing an exception if it is untyped.

Return type:

PsNumericType

class pystencils.backend.functions.GpuFpIntrinsics(value)#

GPU floating point intrinsics.

dividef = ('dividef', 2)#

Fast approximate division

SqrtRn = ('sqrt_rn', 1)#

Fast square root in round-to-nearest-even mode

RSqrtRn = ('rsqrt_rn', 1)#

Fast reciprocal square root in round-to-nearest-even mode

class pystencils.backend.functions.PsGpuIntrinsicFunction(intrin)#

GPU floating point intrinsics

Parameters:

intrin (GpuFpIntrinsics)

class pystencils.backend.functions.GpuGridDimension(value)#

An enumeration.

class pystencils.backend.functions.GpuGridScope(value)#

An enumeration.

class pystencils.backend.functions.PsGpuIndexingFunction(scope, dimension)#

Gpu block, thread, and grid indexing functions.

Calls to IR GPU indexing functions will always typify to the context’s index data type. Platforms must insert appropriate type casts when materializing.

Parameters:
class pystencils.backend.functions.RngSpec(rng_name, dtype, num_ctrs, num_keys, int_arg_type)#

Random number generator specifications for PsRngEngineFunction.

PhiloxFp32 = ('philox_fp32', PsArrayType(element_type=PsIeeeFloatType( width=32, const=False ), shape=(4,), const=False), 4, 2, PsIntegerType( width=32, signed=False, const=False ))#

Philox W=32, N=4 RNG engine returning four float32-values

PhiloxFp64 = ('philox_fp64', PsArrayType(element_type=PsIeeeFloatType( width=64, const=False ), shape=(2,), const=False), 4, 2, PsIntegerType( width=32, signed=False, const=False ))#

Philox W=32, N=4 RNG engine returning two float64-values

class pystencils.backend.functions.PsRngEngineFunction(rng_spec)#

IR function that represents the invocation of a random number generation engine.

This is the IR representation of the symbolic random number generators implemented in pystencils.sympyextensions.random. Each symbolic RNG invocation is mapped onto an RNG engine function through get_for_rng.

Each engine is a function with the signature

engine(ctr_0, ..., ctr_n, key_0, ..., key_m) -> Vec< value_type, K >

which takes n+1 counter arguments, and m+1 key arguments, all of which have the same integer data type (the int_arg_type parameter of the RngSpec). It returns a k-vector of random values of type value_type, modelled using an appropriate PsNamedArrayType.

RNG engine functions must be handled by several transformers, including:

Parameters:

rng_spec (RngSpec) – Specification defining the RNG’s properties

static get_for_rng(state)#

Retrieve the function to be invoked for the given symbolic RNG.

Return type:

PsRngEngineFunction

Parameters:

state (RngState)

class pystencils.backend.functions.CFunction(name, param_types, return_type)#

A concrete C function.

Instances of this class represent a C function by its name, parameter types, and return type.

Parameters:

  • name (str) – Function name

  • param_types (Sequence[PsType]) – Types of the function parameters

  • return_type (PsType) – The function’s return type

static parse(obj)#

Parse the signature of a Python callable object to obtain a CFunction object.

The callable must be fully annotated with type-like objects convertible by create_type.

Return type:

CFunction

Contents