Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Static Transformation properties (manifest) #154

Merged
merged 9 commits into from
Dec 5, 2023
15 changes: 8 additions & 7 deletions loki/bulk/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -614,7 +614,7 @@ def item_successors(self, item):
successors += [self.item_map[child.name]] + self.item_successors(child)
return successors

def process(self, transformation, reverse=False, item_filter=SubroutineItem, use_file_graph=False):
def process(self, transformation):
"""
Process all :attr:`items` in the scheduler's graph

Expand All @@ -638,16 +638,16 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
log = f'[Loki::Scheduler] Applied transformation <{trafo_name}>' + ' in {:.2f}s'
with Timer(logger=info, text=log):

if use_file_graph:
graph = self.file_graph
else:
graph = self.item_graph
# Extract the graph iteration properties from the transformation
graph = self.file_graph if transformation.traverse_file_graph else self.item_graph
item_filter = as_tuple(transformation.item_filter)

# Construct the actual graph to traverse
traversal = nx.topological_sort(graph)
if reverse:
if transformation.reverse_traversal:
traversal = reversed(list(traversal))

if use_file_graph:
if transformation.traverse_file_graph:
for node in traversal:
items = graph.nodes[node]['items']

Expand All @@ -656,6 +656,7 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
if not items:
continue

_item = items[0]
transformation.apply(items[0].source, items=items)
else:
for item in traversal:
Expand Down
5 changes: 4 additions & 1 deletion loki/lint/linter.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,9 @@ class LinterTransformation(Transformation):

_key = 'LinterTransformation'

# This transformation is applied over the file graph
traverse_file_graph = True

def __init__(self, linter, key=None, **kwargs):
self.linter = linter
self.counter = 0
Expand All @@ -272,7 +275,7 @@ def lint_files_scheduler(linter, basedir, config):
"""
scheduler = Scheduler(paths=[basedir], config=SchedulerConfig.from_dict(config))
transformation = LinterTransformation(linter=linter)
scheduler.process(transformation=transformation, use_file_graph=True)
scheduler.process(transformation=transformation)
return transformation.counter


Expand Down
24 changes: 23 additions & 1 deletion loki/transform/build_system_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

from loki.logging import info
from loki.transform.transformation import Transformation
from loki.bulk.item import SubroutineItem, GlobalVarImportItem

__all__ = ['CMakePlanner', 'FileWriteTransformation']

Expand Down Expand Up @@ -129,12 +130,33 @@ class FileWriteTransformation(Transformation):
omitted, it will preserve the original file type.
cuf : bool, optional
Use CUF (CUDA Fortran) backend instead of Fortran backend.
include_module_var_imports : bool, optional
Flag to force the :any:`Scheduler` traversal graph to recognise
module variable imports and write the modified module files.
"""
def __init__(self, builddir=None, mode='loki', suffix=None, cuf=False):

# This transformation is applied over the file graph
traverse_file_graph = True

def __init__(
self, builddir=None, mode='loki', suffix=None, cuf=False,
include_module_var_imports=False
):
self.builddir = Path(builddir)
self.mode = mode
self.suffix = suffix
self.cuf = cuf
self.include_module_var_imports = include_module_var_imports

@property
def item_filter(self):
"""
Override ``item_filter`` to configure whether module variable
imports are honoured in the :any:`Scheduler` traversal.
"""
if self.include_module_var_imports:
return (SubroutineItem, GlobalVarImportItem)
return SubroutineItem

def transform_file(self, sourcefile, **kwargs):
item = kwargs.get('item', None)
Expand Down
21 changes: 9 additions & 12 deletions loki/transform/dependency_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,15 @@ class DependencyTransformation(Transformation):
in the ``ignore``. Default is ``True``.
"""

# This transformation is applied over the file graph
traverse_file_graph = True

# This transformation recurses from the Sourcefile down
recurse_to_modules = True
recurse_to_procedures = True
recurse_to_internal_procedures = False


def __init__(self, suffix, mode='module', module_suffix=None, include_path=None,
replace_ignore_items=True):
self.suffix = suffix
Expand Down Expand Up @@ -151,18 +160,6 @@ def transform_file(self, sourcefile, **kwargs):
if role == 'kernel' and self.mode == 'module':
self.module_wrap(sourcefile, **kwargs)

for module in sourcefile.modules:
# Recursion into contained modules using the sourcefile's "role"
self.transform_module(module, role=role, targets=targets, **kwargs)

if items:
# Recursion into all subroutine items in the current file
for item in items:
self.transform_subroutine(item.routine, item=item, role=item.role, targets=item.targets, **kwargs)
else:
for routine in sourcefile.all_subroutines:
self.transform_subroutine(routine, role=role, targets=targets, **kwargs)

def rename_calls(self, routine, **kwargs):
"""
Update calls to actively transformed subroutines.
Expand Down
15 changes: 9 additions & 6 deletions loki/transform/transform_hoist_variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@
.. code-block:: python

# Transformation: Analysis
scheduler.process(transformation=HoistTemporaryArraysAnalysis(), reverse=True)
scheduler.process(transformation=HoistTemporaryArraysAnalysis())
# Transformation: Synthesis
scheduler.process(transformation=HoistVariablesTransformation())

Expand All @@ -73,10 +73,10 @@
.. code-block:: python

key = "UniqueKey"
scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('b',), key=key), reverse=True)
scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('b',), key=key))
scheduler.process(transformation=HoistTemporaryArraysTransformation(key=key))
key = "AnotherUniqueKey"
scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('a',), key=key), reverse=True)
scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('a',), key=key))
scheduler.process(transformation=HoistTemporaryArraysTransformationAllocatable(key=key))
"""
from loki.expression import FindVariables, SubstituteExpressions
Expand All @@ -100,9 +100,6 @@ class HoistVariablesAnalysis(Transformation):
Create a derived class and override :func:`find_variables<HoistVariablesAnalysis.find_variables>`
to define which variables to be hoisted.

.. note::
To be applied **reversed**, in order to recursively find all variables to be hoisted.

Parameters
----------
key : str
Expand All @@ -112,6 +109,9 @@ class HoistVariablesAnalysis(Transformation):

_key = 'HoistVariablesTransformation'

# Apply in reverse order to recursively find all variables to be hoisted.
reverse_traversal = True

def __init__(self, key=None):
if key is not None:
self._key = key
Expand Down Expand Up @@ -325,6 +325,9 @@ class HoistTemporaryArraysAnalysis(HoistVariablesAnalysis):
for the array dimensions.
"""

# Apply in reverse order to recursively find all variables to be hoisted.
reverse_traversal = True

def __init__(self, key=None, dim_vars=None, **kwargs):
super().__init__(key=key, **kwargs)
self.dim_vars = dim_vars
Expand Down
109 changes: 108 additions & 1 deletion loki/transform/transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from loki.module import Module
from loki.sourcefile import Sourcefile
from loki.subroutine import Subroutine
from loki.bulk.item import SubroutineItem


__all__ = ['Transformation']
Expand All @@ -35,8 +36,53 @@ class Transformation:

Note that in :any:`Sourcefile` objects, all :any:`Module` members will be
traversed before standalone :any:`Subroutine` objects.

Classes inheriting from :any:`Transformation` may configure the
invocation and behaviour during batch processing via a predefined
set of class attributes. These flags determine the underlying
graph traversal when processing complex call trees and determine
how the transformations are invoked for a given type of scheduler
:any:`Item`.

Attributes
----------
reverse_traversal : bool
Forces scheduler traversal in reverse order from the leaf
nodes upwards (default: ``False``).
traverse_file_graph : bool
Apply :any:`Transformation` to the :any:`Sourcefile` object
corresponding to the :any:`Item` being processed, instead of
the program unit in question (default: ``False``).
item_filter : bool
Filter by graph node types to prune the graph and change connectivity.
By default, only calls to :any:`Subroutine` items are used to construct
the graph.
recurse_to_modules : bool
Apply transformation to all :any:`Module` objects when processing
a :any:`Sourcefile` (default ``False``)
recurse_to_procedures : bool
Apply transformation to all :any:`Subroutine` objects when processing
:any:`Sourcefile` or :any:``Module`` objects (default ``False``)
recurse_to_internal_procedures : bool
Apply transformation to all internal :any:`Subroutine` objects
when processing :any:`Subroutine` objects (default ``False``)
"""

# Forces scheduler traversal in reverse order from the leaf nodes upwards
reverse_traversal = False

# Traverse a graph of Sourcefile options corresponding to scheduler items
traverse_file_graph = False
Comment on lines +71 to +75
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we may encounter additional traversal requirements in the future? In which case another variant of specifying these could be to introduce an enum.Flag which has orthogonal traversal options REVERSE_TRAVERSAL, FILE_GRAPH_TRAVERSAL etc. and then allows to combine them in a single traversal attribute on Transformation that is queried by the scheduler. Introducing an additional traversal "mode" would then only require adding the relevant flag.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, in principle I agree, but I think we might want to wait with this until the SGraph is done, as it will introduce more traversal mechanisms. For the sake of it, I gave this a quick try, but it introduces yet more circular dependencies, so I think we might want to postpone for a dedicated PR?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, leave as is and we can pick this up when it actually becomes relevant.


# Filter certain graph nodes to prune the graph and change connectivity
item_filter = SubroutineItem # This can also be a tuple of types

# Recursion behaviour when invoking transformations via ``trafo.apply()``
recurse_to_modules = False # Recurse from Sourcefile to Module
recurse_to_procedures = False # Recurse from Sourcefile/Module to subroutines and functions
recurse_to_internal_procedures = False # Recurse to subroutines in ``contains`` clause


def transform_subroutine(self, routine, **kwargs):
"""
Defines the transformation to apply to :any:`Subroutine` items.
Expand Down Expand Up @@ -120,6 +166,13 @@ def apply_file(self, sourcefile, **kwargs):

This calls :meth:`transform_file`.

If the :attr:`recurse_to_modules` class property is set, it
will also invoke :meth:`apply` on all :any:`Module` objects in
this :any:`Sourcefile`. Likewise, if
:attr:`recurse_to_procedures` is set, it will invoke
:meth:`apply` on all free :any:`Subroutine` objects in this
:any:`Sourcefile`.

Parameters
----------
sourcefile : :any:`Sourcefile`
Expand All @@ -133,15 +186,55 @@ def apply_file(self, sourcefile, **kwargs):
if sourcefile._incomplete:
raise RuntimeError('Transformation.apply_file requires Sourcefile to be complete')

item = kwargs.pop('item', None)
items = kwargs.pop('items', None)
role = kwargs.pop('role', None)
targets = kwargs.pop('targets', None)

if items:
# TODO: This special logic is required for the
# DependencyTransformation to capture certain corner
# cases. Once the module wrapping is split into its
# own transformation, we can probably simplify this.

# We consider the sourcefile to be a "kernel" file if all items are kernels
role = 'kernel' if all(item.role == 'kernel' for item in items) else 'driver'

if targets is None:
# We collect the targets for file/module-level imports from all items
targets = [target for item in items for target in item.targets]

# Apply file-level transformations
self.transform_file(sourcefile, **kwargs)
self.transform_file(sourcefile, item=item, role=role, targets=targets, items=items, **kwargs)

# Recurse to modules, if configured
if self.recurse_to_modules:
for module in sourcefile.modules:
self.transform_module(module, item=item, role=role, targets=targets, items=items, **kwargs)

# Recurse into procedures, if configured
if self.recurse_to_procedures:
if items:
# Recursion into all subroutine items in the current file
for item in items:
self.transform_subroutine(
item.routine, item=item, role=item.role, targets=item.targets, **kwargs
)
else:
for routine in sourcefile.all_subroutines:
self.transform_subroutine(routine, item=item, role=role, targets=targets, **kwargs)

def apply_subroutine(self, subroutine, **kwargs):
"""
Apply transformation to a given :any:`Subroutine` object and its members.

This calls :meth:`transform_subroutine`.

If the :attr:`recurse_to_member_procedures` class property is
set, it will also invoke :meth:`apply` on all
:any:`Subroutine` objects in the ``contains`` clause of this
:any:`Subroutine`.

Parameters
----------
subroutine : :any:`Subroutine`
Expand All @@ -158,12 +251,21 @@ def apply_subroutine(self, subroutine, **kwargs):
# Apply the actual transformation for subroutines
self.transform_subroutine(subroutine, **kwargs)

# Recurse to internal procedures
if self.recurse_to_internal_procedures:
for routine in subroutine.subroutines:
self.apply_subroutine(routine, **kwargs)

def apply_module(self, module, **kwargs):
"""
Apply transformation to a given :any:`Module` object and its members.

This calls :meth:`transform_module`.

If the :attr:`recurse_to_procedures` class property is set,
it will also invoke :meth:`apply` on all :any:`Subroutine`
objects in the ``contains`` clause of this :any:`Module`.

Parameters
----------
module : :any:`Module`
Expand All @@ -180,6 +282,11 @@ def apply_module(self, module, **kwargs):
# Apply the actual transformation for modules
self.transform_module(module, **kwargs)

# Recurse to procedures contained in this module
if self.recurse_to_procedures:
for routine in module.subroutines:
self.apply_subroutine(routine, **kwargs)

def post_apply(self, source, rescope_symbols=False):
"""
Dispatch method for actions to be carried out after applying a transformation
Expand Down
Loading