loki.transformations.replace_kernel

Classes

ReplaceKernels([replace_kernels_map])

Replace configured kernel calls with alternative routines.

class ReplaceKernels(replace_kernels_map=None)

Bases: Transformation

Replace configured kernel calls with alternative routines.

This transformation supports:

• replacing selected call targets by routine name;

• remapping arguments with expressions over the dummy arguments of the original routine (renames, derived-type members, arithmetic, literals);

• passing arguments by keyword or, where possible, by position;

• rewriting USE imports and .intfb.h include-style imports;

• optionally keeping the replacement routine out of scheduler traversal and CMake planning via block=True.

Conceptually, the transformation walks each call site in the processed routine, looks for names listed in replace_kernels_map, resolves the replacement routine, and reconstructs the actual argument list against the replacement signature. The original routine is resolved to match the actual arguments of the call to its dummy arguments, which are then substituted in the rule expressions. Once the replacement call has been rebuilt, the transformation updates imports/includes to match the new callee and optionally marks the replacement routine as blocked for further scheduler traversal.

The replacement map accepts entries of the form:

{
    'old_kernel': {
        'routine': 'new_kernel',
        'args': {
            'new_dummy': 'old_dummy',
            'new_member': 'old_struct%MEMBER',
            'new_end': 'MOD(old_geom%TOTAL%NGPTOT, old_geom%DIM%NPROMA)',
            'new_flag': '.true.',
            'new_start': {'expr': '1', 'positional': True},
        },
        'block': True,
    }
}

Parameters:

replace_kernels_map (dict, optional) –

Case-insensitive mapping from original routine name to replacement configuration. Each configuration entry supports the following keys:

routine

Name of the replacement routine.

args

Optional mapping from replacement dummy argument names to the expression that defines the actual argument of the replacement call. The expression is written in Fortran syntax in terms of the dummy arguments of the original routine, which are substituted by the actual arguments of the original call. Supported forms are:

'new_dummy': 'expression'

Pass the expression by keyword. This covers renaming ('old_dummy'), derived-type members ('old_dummy%MEMBER'), arithmetic and intrinsics, and literals ('.true.').

'new_dummy': {'expr': 'expression', 'positional': True}

Pass the expression by position if all preceding arguments of the replacement call are passed by position, otherwise by keyword.

Replacement dummies without an entry receive the actual argument of the original dummy with the same name. Optional dummies without a match are omitted.

block

If true, mark the replacement routine/module as blocked in the current scheduler item so that it is not traversed or added to the generated plan.

creates_items = True

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.

plan_subroutine(routine, **kwargs)

Define the planning steps to apply for Subroutine items.

For transformations that modify the dependencies of routine (e.g., adding new procedure calls, inlining calls, renaming the interface) this should be implemented. It gets called via the dispatch method apply() if the optional plan_mode argument is set to True.

Parameters:

• routine (Subroutine) – The subroutine to be transformed.

• **kwargs (optional) – Keyword arguments for the transformation.