# 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.