Skip to content

Commit

Permalink
SCC: Let SCCRevector mark routine and SCCAnnotate translates
Browse files Browse the repository at this point in the history
This also renames and refactors the `check_routine_pragmas` utility,
which now only needs to check for genuine `!$loki routine seq`
annotations.
  • Loading branch information
mlange05 committed Aug 9, 2024
1 parent 4a03a00 commit 291ad45
Show file tree
Hide file tree
Showing 7 changed files with 50 additions and 60 deletions.
6 changes: 3 additions & 3 deletions loki/transformations/block_index_transformations.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from loki.transformations.sanitise import resolve_associates
from loki.transformations.utilities import (
recursive_expression_map_update, get_integer_variable,
get_loop_bounds, check_routine_pragmas
get_loop_bounds, check_routine_sequential
)
from loki.transformations.single_column.base import SCCBaseTransformation

Expand Down Expand Up @@ -238,8 +238,8 @@ def process_kernel(self, routine, item, successors, targets, exclude_arrays):
v_index = get_integer_variable(routine, name=self.horizontal.index)
SCCBaseTransformation.resolve_masked_stmts(routine, loop_variable=v_index)

# Bail if routine is marked as sequential or routine has already been processed
if check_routine_pragmas(routine, directive=None):
# Bail if routine is marked as sequential
if check_routine_sequential(routine):
return

bounds = get_loop_bounds(routine, self.horizontal)
Expand Down
21 changes: 13 additions & 8 deletions loki/transformations/single_column/annotate.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from loki.types import DerivedType

from loki.transformations.utilities import (
find_driver_loops, get_local_arrays, check_routine_pragmas
find_driver_loops, get_local_arrays
)


Expand Down Expand Up @@ -116,14 +116,20 @@ def kernel_annotate_sequential_loops_openacc(cls, routine):
@classmethod
def kernel_annotate_subroutine_present_openacc(cls, routine):
"""
Insert ``!$acc data present`` annotations around the body of a subroutine.
Insert ``!$acc routine seq/vector`` directives and wrap
subroutine body in ``!$acc data present`` directives.
Parameters
----------
routine : :any:`Subroutine`
The subroutine to which annotations will be added
"""

# Update `!$loki routine seq/vector` pragmas with `!$acc`
for pragma in FindNodes(ir.Pragma).visit(routine.ir):
if is_loki_pragma(pragma, starts_with='routine'):
pragma._update(keyword='acc')

# Get the names of all array and derived type arguments
args = [a for a in routine.arguments if isinstance(a, sym.Array)]
args += [a for a in routine.arguments if isinstance(a.type.dtype, DerivedType)]
Expand All @@ -146,9 +152,6 @@ def insert_annotations(cls, routine, horizontal):
# to ensure device-resident data is used for array and struct arguments.
cls.kernel_annotate_subroutine_present_openacc(routine)

# Mark routine as `!$acc routine vector` to make it device-callable
routine.spec.append(ir.Pragma(keyword='acc', content='routine vector'))

def transform_subroutine(self, routine, **kwargs):
"""
Apply SCCAnnotate utilities to a :any:`Subroutine`.
Expand Down Expand Up @@ -180,9 +183,11 @@ def process_kernel(self, routine):
Subroutine to apply this transformation to.
"""

# Bail if routine is marked as sequential
if check_routine_pragmas(routine, self.directive):
return
# Bail if this routine has been processed before
for p in FindNodes(ir.Pragma).visit(routine.ir):
# Check if `!$acc routine` has already been added
if p.keyword.lower() == 'acc' and 'routine' in p.content.lower():
return

if self.directive == 'openacc':
self.insert_annotations(routine, self.horizontal)
Expand Down
4 changes: 2 additions & 2 deletions loki/transformations/single_column/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

from loki.transformations.sanitise import resolve_associates
from loki.transformations.utilities import (
get_integer_variable, get_loop_bounds, check_routine_pragmas
get_integer_variable, get_loop_bounds, check_routine_sequential
)


Expand Down Expand Up @@ -146,7 +146,7 @@ def process_kernel(self, routine):
"""

# Bail if routine is marked as sequential or routine has already been processed
if check_routine_pragmas(routine, self.directive):
if check_routine_sequential(routine):
return

# check for horizontal loop bounds in subroutine symbol table
Expand Down
14 changes: 9 additions & 5 deletions loki/transformations/single_column/tests/test_scc.py
Original file line number Diff line number Diff line change
Expand Up @@ -753,9 +753,10 @@ def test_scc_multiple_acc_pragmas(frontend, horizontal, blocking):


@pytest.mark.parametrize('frontend', available_frontends())
def test_scc_base_routine_seq_pragma(frontend, horizontal):
def test_scc_annotate_routine_seq_pragma(frontend, horizontal, blocking):
"""
Test that `!$loki routine seq` pragmas are replaced correctly by `!$acc routine seq` pragmas.
Test that `!$loki routine seq` pragmas are replaced correctly by
`!$acc routine seq` pragmas.
"""

fcode = """
Expand All @@ -781,16 +782,19 @@ def test_scc_base_routine_seq_pragma(frontend, horizontal):
assert pragmas[0].keyword == 'loki'
assert pragmas[0].content == 'routine seq'

transformation = SCCBaseTransformation(horizontal=horizontal, directive='openacc')
transformation.transform_subroutine(routine, role='kernel', targets=['some_kernel',])
transformation = SCCAnnotateTransformation(
horizontal=horizontal, directive='openacc', block_dim=blocking
)
transformation.transform_subroutine(
routine, role='kernel', targets=['some_kernel',]
)

pragmas = FindNodes(Pragma).visit(routine.spec)
assert len(pragmas) == 1
assert pragmas[0].keyword == 'acc'
assert pragmas[0].content == 'routine seq'



@pytest.mark.parametrize('frontend', available_frontends())
def test_scc_vector_reduction(frontend, horizontal, blocking):
"""
Expand Down
11 changes: 9 additions & 2 deletions loki/transformations/single_column/vector.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
from loki.transformations.array_indexing import demote_variables
from loki.transformations.utilities import (
get_integer_variable, get_loop_bounds, find_driver_loops,
get_local_arrays, check_routine_pragmas
get_local_arrays, check_routine_sequential
)


Expand Down Expand Up @@ -99,7 +99,7 @@ def extract_vector_sections(cls, section, horizontal):
# check if calls have been enriched
if not call.routine is BasicType.DEFERRED:
# check if called routine is marked as sequential
if check_routine_pragmas(routine=call.routine, directive=None):
if check_routine_sequential(routine=call.routine):
continue

if call in section:
Expand Down Expand Up @@ -414,6 +414,10 @@ def transform_subroutine(self, routine, **kwargs):
targets = kwargs.get('targets', ())

if role == 'kernel':
# Skip if kernel is marked as `!$loki routine seq`
if check_routine_sequential(routine):
return

# Revector all marked vector sections within the kernel body
routine.body = self.revector_section(routine, routine.body)

Expand All @@ -424,6 +428,9 @@ def transform_subroutine(self, routine, **kwargs):
# Mark sequential loops inside vector sections
self.mark_seq_loops(routine.body)

# Mark subroutine as vector parallel for later annotation
routine.spec.append(ir.Pragma(keyword='loki', content='routine vector'))

if role == 'driver':
with pragmas_attached(routine, ir.Loop):
driver_loops = find_driver_loops(routine=routine, targets=targets)
Expand Down
18 changes: 8 additions & 10 deletions loki/transformations/tests/test_utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
single_variable_declaration, recursive_expression_map_update,
convert_to_lower_case, replace_intrinsics, rename_variables,
get_integer_variable, get_loop_bounds, is_driver_loop,
find_driver_loops, get_local_arrays, check_routine_pragmas
find_driver_loops, get_local_arrays, check_routine_sequential
)


Expand Down Expand Up @@ -511,11 +511,11 @@ def test_transform_utilites_get_local_arrays(frontend, tmp_path):


@pytest.mark.parametrize('frontend', available_frontends())
def test_transform_utilites_check_routine_pragmas(frontend, tmp_path):
""" Test :any:`check_routine_pragmas` utility. """
def test_transform_utilites_check_routine_sequential(frontend, tmp_path):
""" Test :any:`check_routine_sequential` utility. """

fcode = """
module test_check_routine_pragmas_mod
module test_check_routine_sequential_mod
implicit none
contains
Expand All @@ -537,12 +537,10 @@ def test_transform_utilites_check_routine_pragmas(frontend, tmp_path):
i = i + 1
end subroutine test_acc_vec
end module test_check_routine_pragmas_mod
end module test_check_routine_sequential_mod
"""
module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path])

# TODO: This utility needs some serious clean-up, so we're just testing
# the bare basics here and promise to do better next time ;)
assert check_routine_pragmas(module['test_acc_seq'], directive=None)
assert check_routine_pragmas(module['test_loki_seq'], directive=None)
assert check_routine_pragmas(module['test_acc_vec'], directive='openacc')
assert not check_routine_sequential(module['test_acc_seq'])
assert check_routine_sequential(module['test_loki_seq'])
assert not check_routine_sequential(module['test_acc_vec'])
36 changes: 6 additions & 30 deletions loki/transformations/utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
'sanitise_imports', 'replace_selected_kind',
'single_variable_declaration', 'recursive_expression_map_update',
'get_integer_variable', 'get_loop_bounds', 'find_driver_loops',
'get_local_arrays', 'check_routine_pragmas'
'get_local_arrays', 'check_routine_sequential'
]


Expand Down Expand Up @@ -650,41 +650,17 @@ def get_local_arrays(routine, section, unique=True):
return arrays


def check_routine_pragmas(routine, directive):
def check_routine_sequential(routine):
"""
Check if routine is marked as sequential or has already been processed.
Check if routine is marked as "sequential".
Parameters
----------
routine : :any:`Subroutine`
Subroutine to perform checks on.
directive: string or None
Directives flavour to use for parallelism annotations; either
``'openacc'`` or ``None``.
"""

pragmas = FindNodes(ir.Pragma).visit(routine.ir)
routine_pragmas = [p for p in pragmas if p.keyword.lower() in ['loki', 'acc']]
routine_pragmas = [p for p in routine_pragmas if 'routine' in p.content.lower()]

seq_pragmas = [r for r in routine_pragmas if 'seq' in r.content.lower()]
if seq_pragmas:
loki_seq_pragmas = [r for r in routine_pragmas if 'loki' == r.keyword.lower()]
if loki_seq_pragmas:
if directive == 'openacc':
# Mark routine as acc seq
mapper = {seq_pragmas[0]: None}
routine.spec = Transformer(mapper).visit(routine.spec)
routine.body = Transformer(mapper).visit(routine.body)

# Append the acc pragma to routine.spec, regardless of where the corresponding
# loki pragma is found
routine.spec.append(ir.Pragma(keyword='acc', content='routine seq'))
return True

vec_pragmas = [r for r in routine_pragmas if 'vector' in r.content.lower()]
if vec_pragmas:
if directive == 'openacc':
"""
for pragma in FindNodes(ir.Pragma).visit(routine.ir):
if is_loki_pragma(pragma, starts_with='routine seq'):
return True

return False

0 comments on commit 291ad45

Please sign in to comment.