loki.transformations.array_indexing

Collection of utility routines to deal with common array indexing conversions.

Functions

`add_explicit_array_dimensions`(routine)	Make dimensions of arrays explicit within `Subroutine` `routine`.
`demote_variables`(routine, variable_names, ...)	Demote a list of array variables by removing any occurence of a provided set of dimension symbols.
`flatten_arrays`(routine[, order, start_index])	Flatten arrays, converting multi-dimensional arrays to one-dimensional arrays.
`invert_array_indices`(routine)	Invert data/loop accesses from column to row-major
`normalize_array_shape_and_access`(routine)	Shift all arrays to start counting at "1"
`normalize_range_indexing`(routine)	Replace the `(1:size)` indexing in array sizes that OMNI introduces.
`promote_nonmatching_variables`(routine, ...)	Promote multiple variables with potentially non-matching promotion dimensions or index expressions.
`promote_variables`(routine, variable_names, pos)	Promote a list of variables by inserting new array dimensions of given size and updating all uses of these variables with a given index expression.
`promotion_dimensions_from_loop_nest`(...)	Determine promotion dimensions corresponding to the iteration space of a loop nest.
`remove_explicit_array_dimensions`(routine[, ...])	Remove colon notation from array dimensions within `Subroutine` `routine`.
`resolve_vector_notation`(routine)	Resolve implicit vector notation by inserting explicit loops
`shift_to_zero_indexing`(routine[, ignore])	Shift all array indices to adjust to 0-based indexing conventions (eg.

Classes

LowerConstantArrayIndices([...])

A transformation to pass/lower constant array indices down the call tree.

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

shift_to_zero_indexing(routine, ignore=None)

Shift all array indices to adjust to 0-based indexing conventions (eg. for C or Python)

Parameters:

routine (Subroutine) – The subroutine in which the array dimensions should be shifted
ignore (list of str) – List of variable names for which, if found in the dimension expression of an array subscript, that dimension is not shifted to zero.

invert_array_indices(routine)

Invert data/loop accesses from column to row-major

TODO: Take care of the indexing shift between C and Fortran. Basically, we are relying on the CGen to shift the iteration indices and dearly hope that nobody uses the index’s value.

resolve_vector_notation(routine): Resolve implicit vector notation by inserting explicit loops

normalize_range_indexing(routine): Replace the (1:size) indexing in array sizes that OMNI introduces.

promote_variables(routine, variable_names, pos, index=None, size=None)

Promote a list of variables by inserting new array dimensions of given size and updating all uses of these variables with a given index expression.

When providing only size or index, promotion is restricted to updating only variable declarations or their use, respectively, and the other is left unchanged.

Parameters:

routine (Subroutine) – The subroutine in which the variables should be promoted.
variable_names (list of str) – The names of variables to be promoted. Matching of variables against names is case-insensitive.
pos (int) – The position of the new array dimension using Python indexing convention (i.e., count from 0 and use negative values to count from the end).
index (pymbolic.primitives.Expression, optional) – The indexing expression (or a tuple for multi-dimension promotion) to use for the promotion dimension(s), e.g., loop variables. Usage of variables is only updated if index is provided. When the index expression is not live at the variable use, : is used instead.
size (pymbolic.Expression, optional) – The size of the dimension (or tuple for multi-dimension promotion) to insert at pos. When this is provided, the declaration of variables is updated accordingly.

promote_nonmatching_variables(routine, promotion_vars_dims, promotion_vars_index)

Promote multiple variables with potentially non-matching promotion dimensions or index expressions.

This is a convenience routine for using promote_variables() that groups variables by indexing expression and promotion dimensions to reduce the number of calls to promote_variables().

Parameters:

routine (any:Subroutine) – The subroutine to be modified.
promotion_vars_dims (dict) – The mapping of variable names to promotion dimensions. The variables’ shapes are expanded where necessary to have at least these dimensions.
promotion_vars_index (dict) – The mapping of variable names to subscript expressions to be used whenever reading/writing the variable.

promotion_dimensions_from_loop_nest(var_names, loops, promotion_vars_dims, promotion_vars_index)

Determine promotion dimensions corresponding to the iteration space of a loop nest.

Parameters:

var_names (list of str) – The names of variables to consider for promotion.
loops (list of Loop) – The list of nested loops, sorted from outermost to innermost.
promotion_vars_dims (dict((str, tuple))) – The mapping of variable names to promotion dimensions. When determining promotion dimensions for the variables in var_names this dict is checked for already existing promotion dimensions and, if not matching, the maximum of both is taken for each dimension.
promotion_vars_index (dict((str, tuple))) – The mapping of variable names to subscript expressions. These expressions are later inserted for every variable use. When the indexing expression for the loop nest does not match the existing expression in this dict, a RuntimeError is raised.

Returns:

(:data:`promotion_vars_dims`, :data:`promotion_vars_dims`) – The updated mappings promotion_vars_dims and promotion_vars_index.

Return type:

tuple of dict

demote_variables(routine, variable_names, dimensions)

Demote a list of array variables by removing any occurence of a provided set of dimension symbols.

Parameters:

routine (Subroutine) – The subroutine in which the variables should be promoted.
variable_names (list of str) – The names of variables to be promoted. Matching of variables against names is case-insensitive.
dimensions (pymbolic.Expression or tuple) – Symbol name or tuple of symbol names representing the dimension to remove from all occurances of the named variables.

flatten_arrays(routine, order='F', start_index=1)

Flatten arrays, converting multi-dimensional arrays to one-dimensional arrays.

Parameters:

routine (Subroutine) – The subroutine in which the variables should be promoted.
order (str) – Assume Fortran (F) vs. C memory/array order.
start_index (int) – Assume array indexing starts with start_index.

normalize_array_shape_and_access(routine): Shift all arrays to start counting at “1”

class LowerConstantArrayIndices(recurse_to_kernels=True, inline_external_only=True)

Bases: Transformation

A transformation to pass/lower constant array indices down the call tree.

For example, the following code:

subroutine driver(...)
  real, intent(inout) :: var(nlon,nlev,5,nb)
  do ibl=1,10
    call kernel(var(:, :, 1, ibl), var(:, :, 2:5, ibl))
  end do
end subroutine driver

subroutine kernel(var1, var2)
  real, intent(inout) :: var1(nlon, nlev)
  real, intent(inout) :: var2(nlon, nlev, 4)
  var1(:, :) = ...
  do jk=1,nlev
    do jl=1,nlon
      var1(jl, jk) = ...
      do jt=1,4
        var2(jl, jk, jt) = ...
      enddo
    enddo
  enddo
end subroutine kernel

is transformed to:

subroutine driver(...)
  real, intent(inout) :: var(nlon,nlev,5,nb)
  do ibl=1,10
    call kernel(var(:, :, :, ibl), var(:, :, :, ibl))
  end do
end subroutine driver

subroutine kernel(var1, var2)
  real, intent(inout) :: var1(nlon, nlev, 5)
  real, intent(inout) :: var2(nlon, nlev, 5)
  var1(:, :, 1) = ...
  do jk=1,nlev
    do jl=1,nlon
      var1(jl, jk, 1) = ...
      do jt=1,4
        var2(jl, jk, jt + 2 + -1) = ...
      enddo
    enddo
  enddo
end subroutine kernel

Parameters:

recurse_to_kernels (bool) – Recurse to kernels, thus lower constant array indices below the driver level for nested kernel calls (default: True).
inline_external_only (bool) – Inline only external constant expressions or all of them (default: False)

item_filter = (<class 'loki.batch.item.ProcedureItem'>,)

static explicit_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

static is_constant_dim(dim)

Check whether dimension dim is constant, thus, either a constant value or a constant range index.

Parameters:: dim (pymbolic.primitives.Expression)

transform_subroutine(routine, **kwargs)

Defines the transformation to apply to Subroutine items.

For transformations that modify Subroutine objects, this method should be implemented. It gets called via the dispatch method apply().

Parameters:

routine (Subroutine) – The subroutine to be transformed.
**kwargs (optional) – Keyword arguments for the transformation.

process(routine, targets): Process the driver and possibly kernels

update_call_signature(call): Replace constant indices for call arguments being arrays with ‘:’ and update the call.

create_offset_map(call)

Create map/dictionary for arguments with constant array indices.

For, e.g.,

integer :: arg(len1, len2, len3, len4) call kernel(…, arg(:, 2, 4:6, i), …)

offset_map[arg] = {: 0: (0, None, None), # same index as before, no offset 1: (None, 1, len2), # New index, offset 1, size of the dimension is len2 2: (1, 4, len3), # Used to be position 1, offset 4, size of the dimension is len3 3: (-1, None, None), # disregard as this is neither constant nor passed to callee

}

process_callee(routine, offset_map)

Process/adapt the callee according to information in offset_map.

Adapt the variable declarations and usage/indexing.