Introduce 'PragmaModelTransformation' and preliminary OpenMP offload …

…support

loki/batch/tests/test_scheduler.py

@@ -2906,15 +2906,16 @@ def test_pipeline_config_compose(config):
     assert isinstance(pipeline, Pipeline)
 
     # Check that the pipeline is correctly composed
-    assert len(pipeline.transformations) == 8
+    assert len(pipeline.transformations) == 9
     assert type(pipeline.transformations[0]).__name__ == 'RemoveCodeTransformation'
     assert type(pipeline.transformations[1]).__name__ == 'SCCFuseVerticalLoops'
     assert type(pipeline.transformations[2]).__name__ == 'SCCBaseTransformation'
     assert type(pipeline.transformations[3]).__name__ == 'SCCDevectorTransformation'
     assert type(pipeline.transformations[4]).__name__ == 'SCCDemoteTransformation'
     assert type(pipeline.transformations[5]).__name__ == 'SCCVecRevectorTransformation'
     assert type(pipeline.transformations[6]).__name__ == 'SCCAnnotateTransformation'
-    assert type(pipeline.transformations[7]).__name__ == 'ModuleWrapTransformation'
+    assert type(pipeline.transformations[7]).__name__ == 'PragmaModelTransformation'
+    assert type(pipeline.transformations[8]).__name__ == 'ModuleWrapTransformation'
 
     # Check for some specified and default constructor flags
     assert pipeline.transformations[0].call_names == ('dr_hook',)
@@ -2924,7 +2925,7 @@ def test_pipeline_config_compose(config):
     assert pipeline.transformations[2].horizontal.index == 'JL'
     assert pipeline.transformations[2].directive == 'openacc'
     assert pipeline.transformations[3].trim_vector_sections is True
-    assert pipeline.transformations[7].replace_ignore_items is True
+    assert pipeline.transformations[8].replace_ignore_items is True
 
 @pytest.mark.parametrize('frontend', available_frontends())
 @pytest.mark.parametrize('enable_imports', [False, True])

loki/ir/pragma_utils.py

@@ -22,7 +22,8 @@
     'is_loki_pragma', 'get_pragma_parameters', 'process_dimension_pragmas',
     'attach_pragmas', 'detach_pragmas',
     'pragmas_attached', 'attach_pragma_regions', 'detach_pragma_regions',
-    'pragma_regions_attached', 'PragmaAttacher', 'PragmaDetacher'
+    'pragma_regions_attached', 'PragmaAttacher', 'PragmaDetacher',
+    'get_pragma_start_and_parameters'
 ]
 
 
@@ -123,6 +124,40 @@ def _match_spans(open_, close_):
         parameters = {k: v if len(v) > 1 else v[0] for k, v in parameters.items()}
         return parameters
 
+def get_pragma_start_and_parameters(pragma):
+    """
+    Parse a pragma in the form
+
+    ``!$loki command [param] [param_with_val(val)]``
+
+    and return ``command, {param: None, param_with_val: val}``
+
+    Parameters
+    ----------
+    pragma : :any:`Pragma`
+        the pragma to parse and return parameters
+
+    Returns
+    -------
+    tuple(str, dict) :
+        ...
+    """
+    pragma_parameters = PragmaParameters()
+    parameters = defaultdict(list)
+    if pragma.keyword.lower() != 'loki':
+        return None, None
+    content = pragma.content or ''
+    # Remove any line-continuation markers
+    content = content.replace('&', '')
+    starts_with = content.split(' ')[0]
+    if not starts_with:
+        return None, None
+    content = content[len(starts_with):]
+    parameter = pragma_parameters.find(content)
+    for key in parameter:
+        parameters[key].append(parameter[key])
+    parameters = {k: v if len(v) > 1 else v[0] for k, v in parameters.items()}
+    return starts_with, parameters
 
 def get_pragma_parameters(pragma, starts_with=None, only_loki_pragmas=True):
     """
@@ -168,7 +203,6 @@ def get_pragma_parameters(pragma, starts_with=None, only_loki_pragmas=True):
     parameters = {k: v if len(v) > 1 else v[0] for k, v in parameters.items()}
     return parameters
 
-
 def process_dimension_pragmas(ir, scope=None):
     """
     Process any ``!$loki dimension`` pragmas to override deferred dimensions

loki/transformations/__init__.py

@@ -39,3 +39,4 @@
 from loki.transformations.routine_signatures import * # noqa
 from loki.transformations.parallel import * # noqa
 from loki.transformations.dependency import * # noqa
+from loki.transformations.pragma_model import * # noqa

loki/transformations/data_offload/global_var.py

@@ -289,7 +289,7 @@ def transform_module(self, module, **kwargs):
         offload_variables = offload_variables - declared_variables
         if offload_variables:
             module.spec.append(
-                Pragma(keyword='acc', content=f'declare create({", ".join(v.name for v in offload_variables)})')
+                Pragma(keyword='loki', content=f'create device({", ".join(v.name for v in offload_variables)})')
             )
 
     def transform_subroutine(self, routine, **kwargs):
@@ -380,11 +380,11 @@ def process_driver(self, routine, successors):
             copyin_variables = {v for v, _ in uses_symbols if v.parent}
             if update_variables:
                 update_device += (
-                    Pragma(keyword='acc', content=f'update device({", ".join(v.name for v in update_variables)})'),
+                    Pragma(keyword='loki', content=f'update device({", ".join(v.name for v in update_variables)})'),
                 )
             if copyin_variables:
                 update_device += (
-                    Pragma(keyword='acc', content=f'enter data copyin({", ".join(v.name for v in copyin_variables)})'),
+                    Pragma(keyword='loki', content=f'unscoped-data in({", ".join(v.name for v in copyin_variables)})'),
                 )
 
         # All variables that are written in a kernel need a device-to-host transfer
@@ -397,15 +397,17 @@ def process_driver(self, routine, successors):
             }
             if update_variables:
                 update_host += (
-                    Pragma(keyword='acc', content=f'update self({", ".join(v.name for v in update_variables)})'),
+                    Pragma(keyword='loki', content=f'update host({", ".join(v.name for v in update_variables)})'),
                 )
             if copyout_variables:
                 update_host += (
-                    Pragma(keyword='acc', content=f'exit data copyout({", ".join(v.name for v in copyout_variables)})'),
+                    Pragma(keyword='loki',
+                        content=f'unscoped-end-data out({", ".join(v.name for v in copyout_variables)})'),
                 )
             if create_variables:
                 update_device += (
-                    Pragma(keyword='acc', content=f'enter data create({", ".join(v.name for v in create_variables)})'),
+                    Pragma(keyword='loki',
+                        content=f'unscoped-data create({", ".join(v.name for v in create_variables)})'),
                 )
 
         # Replace Loki pragmas with acc data/update pragmas

loki/transformations/data_offload/offload.py

@@ -153,23 +153,25 @@ def insert_data_offload_pragmas(self, routine, targets):
                     if self.assume_deviceptr:
                         offload_args = inargs + outargs + inoutargs
                         if offload_args:
-                            deviceptr = f' deviceptr({", ".join(offload_args)})'
+                            deviceptr = f' vars({", ".join(offload_args)})'
                         else:
                             deviceptr = ''
-                        pragma = Pragma(keyword='acc', content=f'data{deviceptr}')
+                        pragma = Pragma(keyword='loki', content=f'device-ptr{deviceptr}')
+                        pragma_post = Pragma(keyword='loki', content=f'end-device-ptr{deviceptr}')
                     else:
                         offload_args = inargs + outargs + inoutargs
                         if offload_args:
-                            present = f' present({", ".join(offload_args)})'
+                            present = f' vars({", ".join(offload_args)})'
                         else:
                             present = ''
-                        pragma = Pragma(keyword='acc', content=f'data{present}')
+                        pragma = Pragma(keyword='loki', content=f'device-present{present}')
+                        pragma_post = Pragma(keyword='loki', content='end-device-present {present}')
                 else:
-                    copyin = f'copyin({", ".join(inargs)})' if inargs else ''
-                    copy = f'copy({", ".join(inoutargs)})' if inoutargs else ''
-                    copyout = f'copyout({", ".join(outargs)})' if outargs else ''
-                    pragma = Pragma(keyword='acc', content=f'data {copyin} {copy} {copyout}')
-                pragma_post = Pragma(keyword='acc', content='end data')
+                    copyin = f'in({", ".join(inargs)})' if inargs else ''
+                    copy = f'inout({", ".join(inoutargs)})' if inoutargs else ''
+                    copyout = f'out({", ".join(outargs)})' if outargs else ''
+                    pragma = Pragma(keyword='loki', content=f'scoped-data {copyin} {copy} {copyout}')
+                    pragma_post = Pragma(keyword='loki', content='scoped-end-data {copyin} {copy} {copyout}')
                 pragma_map[region.pragma] = pragma
                 pragma_map[region.pragma_post] = pragma_post
 

loki/transformations/data_offload/tests/test_global_var.py

@@ -14,7 +14,7 @@
 
 from loki.transformations import (
     GlobalVariableAnalysis, GlobalVarOffloadTransformation,
-    GlobalVarHoistTransformation
+    GlobalVarHoistTransformation, PragmaModelTransformation
 )
 
 
@@ -259,6 +259,7 @@ def test_global_variable_offload(frontend, key, config, global_variable_analysis
     )
     scheduler.process(GlobalVariableAnalysis(key=key))
     scheduler.process(GlobalVarOffloadTransformation(key=key))
+    scheduler.process(PragmaModelTransformation(directive='openacc'))
     driver = scheduler['#driver'].ir
 
     if key is None:
@@ -346,6 +347,7 @@ def test_transformation_global_var_import(here, config, frontend, tmp_path):
     scheduler = Scheduler(paths=here/'sources/projGlobalVarImports', config=config, frontend=frontend, xmods=[tmp_path])
     scheduler.process(transformation=GlobalVariableAnalysis())
     scheduler.process(transformation=GlobalVarOffloadTransformation())
+    scheduler.process(PragmaModelTransformation(directive='openacc'))
 
     driver = scheduler['#driver'].ir
     moduleA = scheduler['modulea'].ir
@@ -412,6 +414,7 @@ def test_transformation_global_var_import_derived_type(here, config, frontend, t
     scheduler = Scheduler(paths=here/'sources/projGlobalVarImports', config=config, frontend=frontend, xmods=[tmp_path])
     scheduler.process(transformation=GlobalVariableAnalysis())
     scheduler.process(transformation=GlobalVarOffloadTransformation())
+    scheduler.process(PragmaModelTransformation(directive='openacc'))
 
     driver = scheduler['#driver_derived_type'].ir
     module = scheduler['module_derived_type'].ir
@@ -474,6 +477,7 @@ def test_transformation_global_var_hoist(here, config, frontend, hoist_parameter
     scheduler.process(transformation=GlobalVariableAnalysis())
     scheduler.process(transformation=GlobalVarHoistTransformation(hoist_parameters=hoist_parameters,
         ignore_modules=ignore_modules))
+    scheduler.process(PragmaModelTransformation())
 
     driver = scheduler['#driver'].ir
     kernel0 = scheduler['#kernel0'].ir
@@ -576,6 +580,7 @@ def test_transformation_global_var_derived_type_hoist(here, config, frontend, ho
     scheduler = Scheduler(paths=here/'sources/projGlobalVarImports', config=config, frontend=frontend, xmods=[tmp_path])
     scheduler.process(transformation=GlobalVariableAnalysis())
     scheduler.process(transformation=GlobalVarHoistTransformation(hoist_parameters))
+    scheduler.process(PragmaModelTransformation())
 
     driver = scheduler['#driver_derived_type'].ir
     kernel = scheduler['#kernel_derived_type'].ir

loki/transformations/data_offload/tests/test_offload.py

@@ -15,7 +15,7 @@
     pragma_regions_attached, get_pragma_parameters
 )
 
-from loki.transformations import DataOffloadTransformation
+from loki.transformations import DataOffloadTransformation, PragmaModelTransformation
 
 
 @pytest.mark.parametrize('frontend', available_frontends())
@@ -65,16 +65,18 @@ def test_data_offload_region_openacc(caplog, frontend, assume_deviceptr, present
         caplog.clear()
         with caplog.at_level(log_levels['ERROR']):
             with pytest.raises(RuntimeError):
-                data_offload_trafo = DataOffloadTransformation(assume_deviceptr=assume_deviceptr,
-                                                               present_on_device=present_on_device)
+                DataOffloadTransformation(assume_deviceptr=assume_deviceptr, present_on_device=present_on_device)
                 assert len(caplog.records) == 1
                 assert ("[Loki] Data offload: Can't assume device pointer arrays without arrays being marked" +
                     "present on device.") in caplog.records[0].message
             return
 
-    data_offload_trafo = DataOffloadTransformation(assume_deviceptr=assume_deviceptr,
-                                                   present_on_device=present_on_device)
-    driver.apply(data_offload_trafo, role='driver', targets=['kernel_routine'])
+    trafos = ()
+    trafos += (DataOffloadTransformation(assume_deviceptr=assume_deviceptr,
+                                                   present_on_device=present_on_device),)
+    trafos += (PragmaModelTransformation(directive='openacc'),)
+    for trafo in trafos:
+        driver.apply(trafo, role='driver', targets=['kernel_routine'])
 
     pragmas = FindNodes(Pragma).visit(driver.body)
     assert len(pragmas) == 2
@@ -151,8 +153,11 @@ def test_data_offload_region_complex_remove_openmp(frontend):
     kernel = Sourcefile.from_source(fcode_kernel, frontend=frontend)['kernel_routine']
     driver.enrich(kernel)
 
-    offload_transform = DataOffloadTransformation(remove_openmp=True)
-    driver.apply(offload_transform, role='driver', targets=['kernel_routine'])
+    trafos = ()
+    trafos += (DataOffloadTransformation(remove_openmp=True),)
+    trafos += (PragmaModelTransformation(directive='openacc'),)
+    for trafo in trafos:
+        driver.apply(trafo, role='driver', targets=['kernel_routine'])
 
     assert len(FindNodes(Pragma).visit(driver.body)) == 2
     assert all(p.keyword == 'acc' for p in FindNodes(Pragma).visit(driver.body))
@@ -222,7 +227,11 @@ def test_data_offload_region_multiple(frontend):
     kernel = Sourcefile.from_source(fcode_kernel, frontend=frontend)['kernel_routine']
     driver.enrich(kernel)
 
-    driver.apply(DataOffloadTransformation(), role='driver', targets=['kernel_routine'])
+    trafos = ()
+    trafos += (DataOffloadTransformation(),)
+    trafos += (PragmaModelTransformation(directive='openacc'),)
+    for trafo in trafos:
+        driver.apply(trafo, role='driver', targets=['kernel_routine'])
 
     assert len(FindNodes(Pragma).visit(driver.body)) == 2
     assert all(p.keyword == 'acc' for p in FindNodes(Pragma).visit(driver.body))

loki/transformations/pool_allocator.py

@@ -443,14 +443,13 @@ def _get_stack_storage_and_size_var(self, routine, stack_size):
             stack_dealloc = Deallocation(variables=(stack_storage.clone(dimensions=None),))  # pylint: disable=no-member
 
             body_prepend += [stack_alloc]
-            if self.directive == 'openacc':
-                pragma_data_start = Pragma(
-                    keyword='acc',
-                    content=f'data create({stack_storage.name})' # pylint: disable=no-member
-                )
-                body_prepend += [pragma_data_start]
-                pragma_data_end = Pragma(keyword='acc', content='end data')
-                body_append += [pragma_data_end]
+            pragma_data_start = Pragma(
+                keyword='loki',
+                content=f'scoped-data create({stack_storage.name})' # pylint: disable=no-member
+            )
+            body_prepend += [pragma_data_start]
+            pragma_data_end = Pragma(keyword='loki', content=f'scoped-end-data create({stack_storage.name})') # pylint: disable=no-member
+            body_append += [pragma_data_end]
             body_append += [stack_dealloc]
 
         # Inject new variables and body nodes
@@ -759,6 +758,7 @@ def create_pool_allocator(self, routine, stack_size):
         stack_ptr = self._get_stack_ptr(routine)
         stack_end = self._get_stack_end(routine)
 
+        # TODO: generalise using generic Loki pragmas?
         pragma_map = {}
         if self.directive == 'openacc':
             # Find OpenACC loop statements