loki.transformations.single_column.base

Classes

SCCBaseTransformation(horizontal)

A basic set of utilities used in the SCC transformation.
class SCCBaseTransformation(horizontal)

Bases: Transformation

A basic set of utilities used in the SCC transformation. These utilities can either be used as a transformation in their own right, or the contained class methods can be called directly.

Parameters:

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

static rename_index_aliases(routine, dimension)

Rename index aliases: map all index aliases dimension.indices to dimension.index.

Parameters:

  • routine (Subroutine) – The subroutine to rename index aliases.

  • horizontal (Dimension) – Dimension object to rename the index aliases to the first/former index.

static is_elemental(routine)

Check whether Subroutine routine is an elemental routine. Need for distinguishing elemental and non-elemental function to transform those in a different way.

Parameters:

routine (Subroutine) – The subroutine to check whether elemental

transform_subroutine(routine, **kwargs)

Apply SCCBase utilities to a Subroutine.

Notes

The per-routine item config key resolve_vector_notation (bool, default True) can be set to False to skip vector notation resolution for a specific routine. Example scheduler config:

[routines.my_routine]
role = "kernel"
resolve_vector_notation = false
Parameters:

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

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

process_kernel(routine, rename_indices=False, do_resolve_vector_notation=True)

Applies the SCCBase utilities to a “kernel”. This consists of resolving associates, masked statements and vector notation.

Parameters:

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

  • rename_indices (bool, optional) – Whether to rename index aliases to dimension.index. Default is False.

  • do_resolve_vector_notation (bool, optional) – Whether to resolve remaining vector notation (beyond the horizontal dimension) into explicit loops via resolve_vector_notation. The horizontal dimension is always resolved by resolve_vector_dimension regardless of this flag, so that downstream SCC transformations find scalar loop indices instead of range notation. Default is True.

process_driver(routine, do_resolve_vector_notation=True)

Applies the SCCBase utilities to a “driver”. This consists of resolving associates and, by default, vector notation.

Parameters:

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

  • do_resolve_vector_notation (bool, optional) – Whether to resolve vector notation into explicit loops. Set to False to skip this step for driver routines where generated loop bounds may not be available on device. Default is True.