loki.transformations.array_indexing.promote_local_array

Classes

PromoteLocalArrayTransformation(horizontal)

Transformation to promote threadlocal arrays to vector dimension in kernels.
class PromoteLocalArrayTransformation(horizontal, promote_locals=False, preserve_arrays=None, promote_constant_arrays=False)

Bases: Transformation

Transformation to promote threadlocal arrays to vector dimension in kernels.

Parameters:

horizontal (Dimension) – Dimension object describing the variable conventions used in code to define the horizontal data dimension and iteration space.

classmethod get_locals_to_promote(routine, sections, horizontal, promote_constant_arrays=False)

Create a list of local arrays to promote.

Local arrays get promoted if they: * Are used within at least one vector section * Do not already contain the horizontal dimension * Have at least one non-constant dimension (unless

promote_constant_arrays is True)

  • Do not have an entirely deferred shape (all : dimensions), since the actual shape is not known at compile time for pointer/allocatable locals

Parameters:

  • routine (Subroutine) – The subroutine in which to find local arrays.

  • sections (list of Section) – The vector-section IR nodes.

  • horizontal (Dimension) – The horizontal dimension descriptor.

  • promote_constant_arrays (bool, optional) – If True, also promote arrays whose dimensions are all compile-time constants; default: False.

classmethod promote_variables_vector_sections(routine, variable_names, pos, index, sections)

Promote variable uses by inserting a new index dimension, scoped to the given vector sections.

This mirrors the index-insertion logic in promote_variables but instead of using dataflow liveness to decide the index expression per node, it uses vector-section membership: uses inside the given sections get the provided index; remaining uses get :.

Only the variable uses in routine.body are updated; the declarations are not touched. Call promote_variable_declarations separately to update the spec.

Parameters:

  • routine (Subroutine) – The subroutine whose body should be updated.

  • variable_names (list of str) – The names of variables to promote. Matching is case-insensitive.

  • pos (int) – The position of the new array dimension using Python indexing convention (count from 0, negative counts from the end).

  • index (pymbolic.primitives.Expression) – The indexing expression (or tuple for multi-dimension) to use inside vector sections, e.g. the horizontal loop variable.

  • sections (list of Section) – The vector-section IR nodes that define the “inside” region.

transform_subroutine(routine, **kwargs)

Apply Promote utilities to a Subroutine.

Notes

Local array promotion is disabled by default. It can be enabled via the item config key promote_locals.

Parameters:

  • routine (Subroutine) – Subroutine to apply this transformation to.

  • role (string) – Role of the subroutine in the call tree; should be "kernel"

process_kernel(routine, preserve_arrays=None, promote_constant_arrays=False)

Applies the Promote utilities to a “kernel” and promotes all suitable local arrays.

This does two things:

  1. Promotes the declarations of local arrays by appending a new horizontal dimension (using promote_variable_declarations).

  2. Inserts the horizontal index variable into every use of those arrays that appears inside a vector section (using promote_variables_vector_sections).

Parameters:

  • routine (Subroutine) – Subroutine to apply this transformation to.

  • preserve_arrays (list of str, optional) – Names of arrays to exclude from promotion.

  • promote_constant_arrays (bool, optional) – If True, also promote arrays whose dimensions are all compile-time constants; default: False.