Transformation: Add ExtractTransformation and tests

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)

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)