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 by direct rename;

• remapping arguments by extracting a member from a derived-type actual;

• remapping arguments by formatting a custom expression template;

• overriding an argument with a positional value;

• overriding an argument with a literal or expression string;

• 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. When the original call target is unresolved in the IR, the transformation also tries to resolve the original routine from source so that remapping can still be done against the original dummy argument names instead of the raw actual 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': {
            'old_dummy': 'new_dummy',
            'struct_arg': {'map_to': 'new_dummy', 'member': 'MEMBER'},
            'old_end': {
                'map_to': 'new_end',
                'expr': 'MOD({geom}%TOTAL%NGPTOT, {geom}%DIM%NPROMA)',
                'placeholders': {'geom': 'geom'},
            },
            'old_start': {'position': 1},
            'new_flag': '.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 that controls how actual arguments are rebuilt for the replacement call. Supported forms are:

'old_dummy': 'new_dummy'

Map the actual argument of old_dummy onto the replacement dummy new_dummy.

'old_dummy': {'position': 1}

Pass a positional literal or expression to the replacement dummy named old_dummy.

'old_dummy': {'map_to': 'new_dummy', 'member': 'MEMBER'}

Pass actual%MEMBER from the original argument old_dummy to the replacement dummy new_dummy.

'old_dummy': {'map_to': 'new_dummy', 'expr': '...', 'placeholders': {'name': 'old_dummy_name'}}

Build an expression by formatting the template string with the resolved original actual arguments and pass the result to the replacement dummy new_dummy.

'new_dummy': '.true.'

Pass a literal or expression string directly to the replacement dummy new_dummy.

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.