From 5885e78eb706607ae85aa5c03cc566c8630a6376 Mon Sep 17 00:00:00 2001 From: "oakley.brunt" Date: Thu, 19 Sep 2024 11:43:16 +0100 Subject: [PATCH] #2711 dof access added --- .../domain/lfric/kern_call_arg_list.py | 21 +++++++-- src/psyclone/domain/lfric/lfric_loop.py | 36 +++++++--------- .../tests/domain/lfric/dofkern_test.py | 43 ++++++++++++++++++- 3 files changed, 74 insertions(+), 26 deletions(-) diff --git a/src/psyclone/domain/lfric/kern_call_arg_list.py b/src/psyclone/domain/lfric/kern_call_arg_list.py index f451c2a191..7a1ace1af4 100644 --- a/src/psyclone/domain/lfric/kern_call_arg_list.py +++ b/src/psyclone/domain/lfric/kern_call_arg_list.py @@ -379,11 +379,20 @@ def field(self, arg, var_accesses=None): # Look-up the name of the variable that stores the reference to # the data in this field. sym = self._symtab.lookup_with_tag(f"{arg.name}:{suffix}") - # Add the field data array as being read. - self.append(sym.name, var_accesses, var_access_name=sym.name, - mode=arg.access, metadata_posn=arg.metadata_index) - self.psyir_append(Reference(sym)) + if self._kern.is_dofkern: + # If dof kernel, add access to the field by dof ref + dof_sym = self._symtab.find_or_create_integer_symbol( + "df", tag="dof_loop_idx") + dof_sym = Reference(dof_sym) + self.append_array_reference(sym.name, [dof_sym], + ScalarType.Intrinsic.INTEGER, + symbol=sym) + else: + # Add the field data array as being read. + self.append(sym.name, var_accesses, var_access_name=sym.name, + mode=arg.access, metadata_posn=arg.metadata_index) + self.psyir_append(Reference(sym)) def stencil_unknown_extent(self, arg, var_accesses=None): '''Add stencil information to the argument list associated with the @@ -620,6 +629,10 @@ def fs_compulsory_field(self, function_space, var_accesses=None): sym = self.append_array_reference(map_name, [":", ":"], ScalarType.Intrinsic.INTEGER) self.append(sym.name, var_accesses, var_access_name=sym.name) + elif self._kern.is_dofkern: + # Dofmaps are not compatible with user-defined DoF kernels, so + # just pass without creating a dofmap. + pass else: # Pass the dofmap for the cell column cell_name, cell_ref = self.cell_ref_name(var_accesses) diff --git a/src/psyclone/domain/lfric/lfric_loop.py b/src/psyclone/domain/lfric/lfric_loop.py index e1e03642d9..1adc3128cd 100644 --- a/src/psyclone/domain/lfric/lfric_loop.py +++ b/src/psyclone/domain/lfric/lfric_loop.py @@ -225,7 +225,7 @@ def load(self, kern): # Loop bounds self.set_lower_bound("start") const = LFRicConstants() - if isinstance(kern, LFRicBuiltIn): + if isinstance(kern, LFRicBuiltIn) or kern.is_dofkern: # If the kernel is a built-in/pointwise operation # then this loop must be over DoFs if Config.get().api_conf("lfric").compute_annexed_dofs \ @@ -234,10 +234,6 @@ def load(self, kern): self.set_upper_bound("nannexed") else: self.set_upper_bound("ndofs") - elif kern.is_dofkern: - # If the kernel is user-defined and operats on dofs, always - # set upper bound to 'ndofs' - self.set_upper_bound("ndofs") else: if Config.get().distributed_memory: if self._field.is_operator: @@ -477,22 +473,22 @@ def _upper_bound_fortran(self): sym = sym_tab.find_or_create_tag(root_name) return f"{sym.name}(colour, {depth})" if self._upper_bound_name in ["ndofs", "nannexed"]: - if isinstance(self._kern, LFRicBuiltIn): - if Config.get().distributed_memory: - if self._upper_bound_name == "ndofs": - result = ( - f"{self.field.proxy_name_indexed}%" - f"{self.field.ref_name()}%get_last_dof_owned()") - else: # nannexed - result = ( - f"{self.field.proxy_name_indexed}%" - f"{self.field.ref_name()}%get_last_dof_annexed()") - else: - result = self._kern.undf_name + if Config.get().distributed_memory: + if self._upper_bound_name == "ndofs": + result = ( + f"{self.field.proxy_name_indexed}%" + f"{self.field.ref_name()}%get_last_dof_owned()") + else: # nannexed + result = ( + f"{self.field.proxy_name_indexed}%" + f"{self.field.ref_name()}%get_last_dof_annexed()") else: - # User-defined dof kernel has undf_name in a different location - # so use that here instead of calling `get_undf` again - result = self._field_space.undf_name + if isinstance(self._kern, LFRicBuiltIn): + result = self._kern.undf_name + else: + # User-defined dof kernel has undf_name in a different + # location + result = self._field_space.undf_name return result if self._upper_bound_name == "ncells": if Config.get().distributed_memory: diff --git a/src/psyclone/tests/domain/lfric/dofkern_test.py b/src/psyclone/tests/domain/lfric/dofkern_test.py index 1f0dd9fda5..4293f7b21a 100644 --- a/src/psyclone/tests/domain/lfric/dofkern_test.py +++ b/src/psyclone/tests/domain/lfric/dofkern_test.py @@ -37,13 +37,15 @@ This module tests the metadata validation in LFRicKernMetadata of user-supplied kernels operating on degrees of freedom (dofs) ''' - +import os import pytest from fparser import api as fpapi from psyclone.configuration import Config -from psyclone.domain.lfric import LFRicKernMetadata +from psyclone.domain.lfric import LFRicKernMetadata, LFRicKern, LFRicLoop +from psyclone.parse.algorithm import parse +from psyclone.psyGen import PSyFactory from psyclone.parse.utils import ParseError @@ -54,6 +56,9 @@ def setup(): yield Config._instance = None +BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname( + os.path.abspath(__file__)))), "test_files", "dynamo0p3") +TEST_API = "lfric" CODE = ''' module testkern_dofs_mod @@ -142,3 +147,37 @@ def test_dof_kernel_invalid_field_vector(): assert ("Kernel 'testkern_dofs_type' operates on 'dof' but has a vector " "argument 'gh_field*3'. This is not permitted in the LFRic API." in str(excinfo.value)) + + +def test_dof_kernel(): + ''' Check that we raise an exception if we encounter metadata + for a dof kernel with a field vector. + + ''' + # Substitute field for field vector + code = CODE.replace( + """ + (/ arg_type(gh_field, gh_real, gh_write, w1), & + arg_type(gh_field, gh_real, gh_read, w2) & + """, + """ + (/ arg_type(gh_field, gh_real, gh_write, w1), & + arg_type(gh_field, gh_real, gh_read, w1) & + """, + 1) + ast = fpapi.parse(code, ignore_comments=False) + name = "testkern_dofs_type" + md = LFRicKernMetadata(ast, name=name) + kern = LFRicKern() + kern.load_meta(ktype=md) + + +def test_gen(): + _, invoke_info = parse(os.path.join(BASE_PATH, + "1.14_single_invoke_dofs.f90"), + api=TEST_API) + print(invoke_info) + psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info) + sched = psy.invokes.invoke_list[0].schedule + print(sched.view()) + print(psy.gen)