loki.transformations.array_indexing.vector_notation

Utilities to manipulate vector notation in array expressions.

Functions

add_explicit_array_dimensions(routine)

Make dimensions of arrays explicit within Subroutine routine.

remove_explicit_array_dimensions(routine[, ...])

Remove colon notation from array dimensions within Subroutine routine.

resolve_vector_dimension(routine, dimension)

Resolve vector notation for a given dimension only.

resolve_vector_notation(routine[, ...])

Resolve implicit vector notation by inserting explicit loops.

Classes

IterationRangeIndexMapper(*args[, loop_map, ...])

A mapper that replaces fully qualified RangeIndex symbols with discrete loop indices and collects the according index_to_range_map.

IterationRangeShapeMapper()

A mapper that derives the fully qualified iteration dimension for unbounded RangeIndex indices in array expressions.

ResolveVectorNotationTransformer(*args[, ...])

A Transformer that resolves implicit vector notation by inserting explicit loops.
remove_explicit_array_dimensions(routine, calls_only=False)

Remove colon notation from array dimensions within Subroutine routine. E.g., convert two-dimensional array arr2d(:,:) to arr2d or arr3d(:,:,:) to arr3d, but NOT e.g., arr(1,:,:).

Parameters:

  • routine (Subroutine) – The subroutine to check

  • calls_only (bool) – Whether to remove colon notation from array dimensions only from arrays within (inline) calls or all arrays (default: False)

add_explicit_array_dimensions(routine)

Make dimensions of arrays explicit within Subroutine routine. E.g., convert two-dimensional array arr2d to arr2d(:,:) or arr3d to arr3d(:,:,:).

Parameters:

routine (Subroutine) – The subroutine to check

resolve_vector_notation(routine, resolve_implicit_rhs_ranges=True, substitute_derived_type_bounds=False)

Resolve implicit vector notation by inserting explicit loops.

Parameters:

  • routine (Subroutine) – The subroutine in which to resolve vector notation usage.

  • resolve_implicit_rhs_ranges (bool) – When True (default), resolve all LHS range dimensions even if the corresponding RHS arrays use bare : ranges.

  • substitute_derived_type_bounds (bool) – When True, replace derived-type member references in synthesized loop bounds with plain scalar variables. Only needed for driver routines. Defaults to False.

resolve_vector_dimension(routine, dimension, derive_qualified_ranges=False, resolve_implicit_rhs_ranges=True, substitute_derived_type_bounds=False)

Resolve vector notation for a given dimension only. The dimension is defined by a loop variable and the bounds of the given range.

Unlike the related resolve_vector_notation() utility, this will only resolve the defined dimension according to bounds and loop_variable.

Parameters:

  • routine (Subroutine) – The subroutine in which to resolve vector notation usage.

  • dimension (Dimension) – Dimension object that defines the dimension to resolve

  • derive_qualified_ranges (bool) – Flag to enable the derivation of (all) range bounds from shape information.

  • resolve_implicit_rhs_ranges (bool) – When True (default), resolve all LHS range dimensions even if the corresponding RHS arrays use bare : ranges.

  • substitute_derived_type_bounds (bool) – When True, replace derived-type member references in synthesized loop bounds with plain scalar variables. Only needed for driver routines. Defaults to False.

class ResolveVectorNotationTransformer(*args, loop_map=None, scope=None, derive_qualified_ranges=True, map_unknown_ranges=True, resolve_implicit_rhs_ranges=True, substitute_derived_type_bounds=False, **kwargs)

Bases: Transformer

A Transformer that resolves implicit vector notation by inserting explicit loops.

Parameters:

  • loop_map (dict of tuple to Variable) – A dict mapping the tuple (lower, upper, step) to a known variable symbol to use as loop index.

  • scope (Subroutine or Module) – The scope in which to create new loop index variables

  • derive_qualified_ranges (bool) – Derive explicit bounds for all unqualified index ranges (:) before resolving them with loops.

  • map_unknown_ranges (bool) – Flag to indicate whether unknown, but fully qualified range indices are to be remapped to loops.

  • resolve_implicit_rhs_ranges (bool) – When True (default), resolve all LHS range dimensions even if the corresponding RHS arrays use bare : (unqualified) ranges. When False, only resolve dimensions where all RHS arrays have explicit (qualified) ranges.

  • substitute_derived_type_bounds (bool) – When True, replace derived-type member references in synthesized loop bounds with existing or newly created plain scalar variables (see _substitute_derived_type_bounds()). This is intended for driver routines where device-safe plain scalars are required. Defaults to False; kernels should leave derived-type bounds as-is.

visit_Assignment(stmt, **kwargs)
visit_MaskedStatement(masked, **kwargs)