From 8be5434f21ce7a9daedb46fd000f187df3fbc2ab Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Sep 2024 18:42:16 +0000 Subject: [PATCH] Transformation: Add `ExtractTransformation` and tests --- loki/transformations/extract/__init__.py | 68 +++++++++ .../tests/test_extract_transformation.py | 133 ++++++++++++++++++ 2 files changed, 201 insertions(+) create mode 100644 loki/transformations/extract/tests/test_extract_transformation.py diff --git a/loki/transformations/extract/__init__.py b/loki/transformations/extract/__init__.py index c6212cc33..6cde05b07 100644 --- a/loki/transformations/extract/__init__.py +++ b/loki/transformations/extract/__init__.py @@ -7,7 +7,75 @@ """ Transformations sub-package that provides various forms of source-code extraction into standalone :any:`Subroutine` objects. + +The various extractions mechanisms are provided as standalone utility +methods, or via the :any:`ExtractTransformation` class for for batch +processing. + +These utilities represent the conceptual inverse operation to +"inlining", as done by the :any:`InlineTransformation`. """ from loki.transformations.extract.internal import * # noqa from loki.transformations.extract.marked import * # noqa + +from loki.batch import Transformation + + +__all__ = ['ExtractTransformation'] + + +class ExtractTransformation(Transformation): + """ + :any:`Transformation` class to apply several types of source + extraction when batch-processing large source trees via the + :any:`Scheduler`. + + Parameters + ---------- + inline_internals : bool + Extract internal procedure (see :any:`extract_internal_procedures`); + default: False. + inline_marked : bool + Extract :any:`Subroutine` objects marked by pragma annotations + (see :any:`extract_marked_subroutines`); default: True. + """ + def __init__(self, extract_internals=False, extract_marked=True): + self.extract_internals = extract_internals + self.extract_marked = extract_marked + + def transform_module(self, module, **kwargs): + """ + Extract internals procedures and marked subroutines and add + them to the given :any:`Module`. + """ + + # Extract internal (contained) procedures into standalone ones + if self.extract_internals: + for routine in module.subroutines: + new_routines = extract_internal_procedures(routine) + module.contains.append(new_routines) + + # Extract pragma-marked code regions into standalone subroutines + if self.extract_marked: + for routine in module.subroutines: + new_routines = extract_marked_subroutines(routine) + module.contains.append(new_routines) + + def transform_file(self, sourcefile, **kwargs): + """ + Extract internals procedures and marked subroutines and add + them to the given :any:`Sourcefile`. + """ + + # Extract internal (contained) procedures into standalone ones + if self.extract_internals: + for routine in sourcefile.subroutines: + new_routines = extract_internal_procedures(routine) + sourcefile.ir.append(new_routines) + + # Extract pragma-marked code regions into standalone subroutines + if self.extract_marked: + for routine in sourcefile.subroutines: + new_routines = extract_marked_subroutines(routine) + sourcefile.ir.append(new_routines) diff --git a/loki/transformations/extract/tests/test_extract_transformation.py b/loki/transformations/extract/tests/test_extract_transformation.py new file mode 100644 index 000000000..8f3364842 --- /dev/null +++ b/loki/transformations/extract/tests/test_extract_transformation.py @@ -0,0 +1,133 @@ +# (C) Copyright 2018- ECMWF. +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +import pytest + +from loki import Subroutine, Module, Sourcefile +from loki.frontend import available_frontends +from loki.ir import nodes as ir, FindNodes + +from loki.transformations.extract import ExtractTransformation + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('extract_marked', [False, True]) +@pytest.mark.parametrize('extract_internals', [False, True]) +def test_extract_transformation_module(extract_internals, extract_marked, frontend): + """ + Test basic subroutine extraction from marker pragmas in modules. + """ + fcode = """ +module test_extract_mod +implicit none +contains + +subroutine outer(n, a, b) + integer, intent(in) :: n + real(kind=8), intent(inout) :: a, b(n) + real(kind=8) :: x(n), y(n, n+1) + integer :: i, j + + x(:) = a + do i=1, n + y(i,:) = b(i) + end do + + !$loki extract name(test1) + do i=1, n + do j=1, n+1 + x(i) = x(i) + y(i, j) + end do + end do + !$loki end extract + + do i=1, n + call plus_one(x, i=i) + end do + +contains + subroutine plus_one(f, i) + real(kind=8), intent(inout) :: f(:) + integer, intent(in) :: i + + f(i) = f(i) + 1.0 + end subroutine plus_one +end subroutine outer +end module test_extract_mod +""" + module = Module.from_source(fcode, frontend=frontend) + + ExtractTransformation( + extract_internals=extract_internals, extract_marked=extract_marked + ).apply(module) + + routines = tuple(r for r in module.contains.body if isinstance(r, Subroutine)) + assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if extract_marked else 0) + assert ('plus_one' in module) == extract_internals + assert ('test1' in module) == extract_marked + + outer = module['outer'] + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if extract_marked else 1) + outer_internals = tuple(r for r in outer.contains.body if isinstance(r, Subroutine)) + assert len(outer_internals) == (0 if extract_internals else 1) + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('extract_marked', [False, True]) +@pytest.mark.parametrize('extract_internals', [False, True]) +def test_extract_transformation_sourcefile(extract_internals, extract_marked, frontend): + """ + Test internal procedure extraction from subroutines. + """ + fcode = """ +subroutine outer(n, a, b) + integer, intent(in) :: n + real(kind=8), intent(inout) :: a, b(n) + real(kind=8) :: x(n), y(n, n+1) + integer :: i, j + + x(:) = a + do i=1, n + y(i,:) = b(i) + end do + + !$loki extract name(test1) + do i=1, n + do j=1, n+1 + x(i) = x(i) + y(i, j) + end do + end do + !$loki end extract + + do i=1, n + call plus_one(x, i=i) + end do + +contains + subroutine plus_one(f, i) + real(kind=8), intent(inout) :: f(:) + integer, intent(in) :: i + + f(i) = f(i) + 1.0 + end subroutine plus_one +end subroutine outer +""" + sourcefile = Sourcefile.from_source(fcode, frontend=frontend) + + ExtractTransformation( + extract_internals=extract_internals, extract_marked=extract_marked + ).apply(sourcefile) + + routines = tuple(r for r in sourcefile.ir.body if isinstance(r, Subroutine)) + assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if extract_marked else 0) + assert ('plus_one' in sourcefile) == extract_internals + assert ('test1' in sourcefile) == extract_marked + + outer = sourcefile['outer'] + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if extract_marked else 1) + outer_internals = tuple(r for r in outer.contains.body if isinstance(r, Subroutine)) + assert len(outer_internals) == (0 if extract_internals else 1)