loki.transformations.array_indexing.array_indices

Utilities to change indices and indexing in array expressions.

Functions

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

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.

normalize_range_indexing(routine)

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

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.