Reference Guide  2.5.0
kernel_interface.py
1 # -----------------------------------------------------------------------------
2 # BSD 3-Clause License
3 #
4 # Copyright (c) 2020-2024, Science and Technology Facilities Council.
5 # All rights reserved.
6 #
7 # Redistribution and use in source and binary forms, with or without
8 # modification, are permitted provided that the following conditions are met:
9 #
10 # * Redistributions of source code must retain the above copyright notice, this
11 # list of conditions and the following disclaimer.
12 #
13 # * Redistributions in binary form must reproduce the above copyright notice,
14 # this list of conditions and the following disclaimer in the documentation
15 # and/or other materials provided with the distribution.
16 #
17 # * Neither the name of the copyright holder nor the names of its
18 # contributors may be used to endorse or promote products derived from
19 # this software without specific prior written permission.
20 #
21 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
22 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
23 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
24 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
25 # COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
26 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
27 # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
28 # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
29 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
30 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
31 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
32 # POSSIBILITY OF SUCH DAMAGE.
33 # -----------------------------------------------------------------------------
34 # Author R. W. Ford, STFC Daresbury Lab
35 # Modified: I. Kavcic and L. Turner, Met Office
36 # A. R. Porter and N. Nobre, STFC Daresbury Lab
37 # Modified: J. Henrichs, Bureau of Meteorology
38 
39 '''This module creates the expected arguments for an LFRic coded
40 kernel based on the kernel metadata.
41 
42 '''
43 from psyclone.core import AccessType
44 from psyclone.domain.lfric import ArgOrdering, LFRicConstants
45 from psyclone.domain.lfric.lfric_symbol_table import LFRicSymbolTable
46 from psyclone.domain.lfric.lfric_types import LFRicTypes
47 from psyclone.errors import InternalError
48 from psyclone.psyir.frontend.fparser2 import INTENT_MAPPING
49 from psyclone.psyir.nodes import Reference
50 from psyclone.psyir.symbols import ArgumentInterface
51 
52 
53 # pylint: disable=too-many-public-methods, no-member
55  '''Create the kernel arguments for the supplied kernel as specified by
56  the associated kernel metadata and the kernel ordering rules
57  encoded in the ArgOrdering base class as method callbacks.
58 
59  A PSyIR symbol table is created containing appropriate LFRic PSyIR
60  symbols to specify the arguments. If an argument is an array with
61  one or more dimension sizes specified by another argument, then
62  the associated array symbol is created so that it references the
63  appropriate symbol.
64 
65  Related arguments - e.g. a field has an associated dofmap - are
66  not directly connected, they must be inferred from the function
67  space names. It is not yet clear whether this would be useful or
68  not.
69 
70  TODO #928: This class should replace the current kernel stub generation
71  code when all of its methods are implemented.
72 
73  :param kern: the kernel for which to create arguments.
74  :type kern: :py:class:`psyclone.domain.lfric.LFRicKern`
75 
76  '''
77  #: Mapping from a generic PSyIR datatype to the equivalent
78  #: LFRic-specific field datasymbol.
79  field_mapping = {
80  "integer": "IntegerFieldDataSymbol",
81  "real": "RealFieldDataSymbol",
82  "logical": "LogicalFieldDataSymbol"}
83  #: Mapping from a generic PSyIR datatype to the equivalent
84  #: LFRic-specific vector field datasymbol.
85  vector_field_mapping = {
86  "integer": "IntegerVectorFieldDataSymbol",
87  "real": "RealVectorFieldDataSymbol",
88  "logical": "LogicalVectorFieldDataSymbol"}
89  #: Mapping from the LFRic metadata description of quadrature to the
90  #: associated LFRic-specific basis function datasymbol.
91  basis_mapping = {
92  "gh_quadrature_xyoz": "BasisFunctionQrXyozDataSymbol",
93  "gh_quadrature_face": "BasisFunctionQrFaceDataSymbol",
94  "gh_quadrature_edge": "BasisFunctionQrEdgeDataSymbol"}
95  #: Mapping from the LFRic metadata description of quadrature to the
96  #: associated LFRic-specific differential basis function datasymbol.
97  diff_basis_mapping = {
98  "gh_quadrature_xyoz": "DiffBasisFunctionQrXyozDataSymbol",
99  "gh_quadrature_face": "DiffBasisFunctionQrFaceDataSymbol",
100  "gh_quadrature_edge": "DiffBasisFunctionQrEdgeDataSymbol"}
101  _read_access = ArgumentInterface(ArgumentInterface.Access.READ)
102 
103  def __init__(self, kern):
104  super().__init__(kern)
105  # We need a brand new symbol table, not a reference to the one for
106  # the Schedule containing this kernel call (which is what the
107  # ArgOrdering constructor defaults to). This is so that we can
108  # specify the correct interface for those symbols passed
109  # as arguments.
110  self._forced_symtab_forced_symtab_forced_symtab = LFRicSymbolTable()
111 
112  def generate(self, var_accesses=None):
113  '''Call the generate base class then add the argument list as it can't
114  be appended as we go along.
115 
116  :param var_accesses: an unused optional argument that stores \
117  information about variable accesses.
118  :type var_accesses: :\
119  py:class:`psyclone.core.VariablesAccessInfo`
120 
121  '''
122  super().generate(var_accesses=var_accesses)
123  # Set the argument list for the symbol table. This is done at
124  # the end after incrementally adding symbols to the _args
125  # list, as it is not possible to incrementally add symbols to
126  # the symbol table argument list.
127  self._symtab_symtab.specify_argument_list(self._arglist_arglist_arglist)
128  # While data dependence analysis does not use the symbol
129  # table, see #845, we have to provide variable information
130  # separately. This is done by using the base class append()
131  # method. However, this method also adds the variable names to the
132  # internal _arglist list which we do not want as we have
133  # already added our symbols there. Therefore we need to remove
134  # them afterwards.
135  # Map from symbol table accesses to dependence analysis accesses.
136  mapping = {ArgumentInterface.Access.READ: AccessType.READ,
137  ArgumentInterface.Access.READWRITE: AccessType.READWRITE,
138  ArgumentInterface.Access.WRITE: AccessType.WRITE}
139  len_arglist = len(self._arglist_arglist_arglist)
140  for symbol in self._symtab_symtab.symbols:
141  self.appendappend(symbol.name, var_accesses,
142  mode=mapping[symbol.interface.access])
143  self._arglist_arglist_arglist = self._arglist_arglist_arglist[:len_arglist]
144 
145  def cell_position(self, var_accesses=None):
146  '''Create an LFRic cell-position object and add it to the symbol table
147  and argument list.
148 
149  :param var_accesses: an unused optional argument that stores \
150  information about variable accesses.
151  :type var_accesses: :\
152  py:class:`psyclone.core.VariablesAccessInfo`
153 
154  '''
155  symbol = self._symtab_symtab.find_or_create_tag(
156  "cell", symbol_type=LFRicTypes("CellPositionDataSymbol"),
157  interface=self._read_access_read_access)
158  self._arglist_arglist_arglist.append(symbol)
159 
160  def mesh_height(self, var_accesses=None):
161  '''Create an LFRic mesh height object and add it to the symbol table
162  and argument list.
163 
164  :param var_accesses: an unused optional argument that stores \
165  information about variable accesses.
166  :type var_accesses: :\
167  py:class:`psyclone.core.VariablesAccessInfo`
168 
169  '''
170  symbol = self._symtab_symtab.find_or_create_tag(
171  "nlayers", symbol_type=LFRicTypes("MeshHeightDataSymbol"),
172  interface=self._read_access_read_access)
173  self._arglist_arglist_arglist.append(symbol)
174 
175  def _mesh_ncell2d(self, var_accesses=None):
176  '''Not implemented.
177 
178  :param var_accesses: an unused optional argument that stores \
179  information about variable accesses.
180  :type var_accesses: :\
181  py:class:`psyclone.core.VariablesAccessInfo`
182 
183  :raises NotImplementedError: as this method is not implemented.
184 
185  '''
186  raise NotImplementedError("TODO #928: _mesh_ncell2d not implemented")
187 
188  def _mesh_ncell2d_no_halos(self, var_accesses=None):
189  '''Not implemented.
190 
191  :param var_accesses: an unused optional argument that stores \
192  information about variable accesses.
193  :type var_accesses: :\
194  py:class:`psyclone.core.VariablesAccessInfo`
195 
196  :raises NotImplementedError: as this method is not implemented.
197 
198  '''
199  raise NotImplementedError(
200  "TODO #928: _mesh_ncell2d_no_halos not implemented")
201 
202  def cell_map(self, var_accesses=None):
203  '''Not implemented.
204 
205  :param var_accesses: an unused optional argument that stores \
206  information about variable accesses.
207  :type var_accesses: :\
208  py:class:`psyclone.core.VariablesAccessInfo`
209 
210  :raises NotImplementedError: as this method is not implemented.
211 
212  '''
213  raise NotImplementedError("TODO #928: cell_map not implemented")
214 
215  def field_vector(self, argvect, var_accesses=None):
216  '''Create LFRic field vector arguments and add them to the symbol
217  table and argument list. Also declare the associated "undf"
218  symbol if it has not already been declared so that it can be
219  used to dimension the field vector arguments.
220 
221  :param argvect: the field vector to add.
222  :type argvect: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
223  :param var_accesses: an unused optional argument that stores \
224  information about variable accesses.
225  :type var_accesses: :\
226  py:class:`psyclone.core.VariablesAccessInfo`
227 
228  :raises NotImplementedError: if the datatype of the vector \
229  field is not supported.
230 
231  '''
232  fs_name = argvect.function_space.orig_name
233  undf_symbol = self._symtab_symtab.find_or_create_tag(
234  f"undf_{fs_name}", fs=fs_name,
235  symbol_type=LFRicTypes("NumberOfUniqueDofsDataSymbol"),
236  interface=self._read_access_read_access)
237 
238  interface = ArgumentInterface(INTENT_MAPPING[argvect.intent])
239  try:
240  type_name = self.vector_field_mappingvector_field_mapping[argvect.intrinsic_type]
241  field_class = LFRicTypes(type_name)
242  except KeyError as info:
243  message = (f"kernel interface does not support a vector field of "
244  f"type '{argvect.intrinsic_type}'.")
245  raise NotImplementedError(message) from info
246  for idx in range(argvect.vector_size):
247  tag = f"{argvect.name}_v{idx}"
248  field_data_symbol = self._symtab_symtab.find_or_create_tag(
249  tag, symbol_type=field_class, dims=[Reference(undf_symbol)],
250  fs=fs_name, interface=interface)
251 
252  self._arg_index_to_metadata_index_arg_index_to_metadata_index[len(self._arglist_arglist_arglist)] = (
253  argvect.metadata_index)
254  self._arglist_arglist_arglist.append(field_data_symbol)
255 
256  def field(self, arg, var_accesses=None):
257  '''Create an LFRic field argument and add it to the symbol table and
258  argument list. Also declare the associated "undf" symbol if it
259  has not already been declared so that it can be used to
260  dimension the field argument.
261 
262  :param arg: the field to add.
263  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
264  :param var_accesses: an unused optional argument that stores \
265  information about variable accesses.
266  :type var_accesses: :\
267  py:class:`psyclone.core.VariablesAccessInfo`
268 
269  :raises NotImplementedError: if the datatype of the field is \
270  not supported.
271 
272  '''
273  fs_name = arg.function_space.orig_name
274 
275  undf_symbol = self._symtab_symtab.find_or_create_tag(
276  f"undf_{fs_name}",
277  symbol_type=LFRicTypes("NumberOfUniqueDofsDataSymbol"),
278  fs=fs_name, interface=self._read_access_read_access)
279 
280  try:
281  type_name = self.field_mappingfield_mapping[arg.intrinsic_type]
282  field_class = LFRicTypes(type_name)
283  except KeyError as info:
284  message = (f"kernel interface does not support a field of type "
285  f"'{arg.intrinsic_type}'.")
286  raise NotImplementedError(message) from info
287  field_data_symbol = self._symtab_symtab.find_or_create_tag(
288  arg.name, interface=ArgumentInterface(INTENT_MAPPING[arg.intent]),
289  symbol_type=field_class, dims=[Reference(undf_symbol)], fs=fs_name)
290 
291  self._arg_index_to_metadata_index_arg_index_to_metadata_index[len(self._arglist_arglist_arglist)] = (
292  arg.metadata_index)
293  self._arglist_arglist_arglist.append(field_data_symbol)
294 
295  def stencil_unknown_extent(self, arg, var_accesses=None):
296  '''Not implemented.
297 
298  :param arg: the kernel argument with which the stencil is associated.
299  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
300  :param var_accesses: an unused optional argument that stores \
301  information about variable accesses.
302  :type var_accesses: :\
303  py:class:`psyclone.core.VariablesAccessInfo`
304 
305  :raises NotImplementedError: as this method is not implemented.
306 
307  '''
308  raise NotImplementedError(
309  "TODO #928: stencil_unknown_extent not implemented")
310 
311  def stencil_unknown_direction(self, arg, var_accesses=None):
312  '''Not implemented.
313 
314  :param arg: the kernel argument with which the stencil is associated.
315  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
316  :param var_accesses: an unused optional argument that stores \
317  information about variable accesses.
318  :type var_accesses: :\
319  py:class:`psyclone.core.VariablesAccessInfo`
320 
321  :raises NotImplementedError: as this method is not implemented.
322 
323  '''
324  raise NotImplementedError(
325  "TODO #928: stencil_unknown_direction not implemented")
326 
327  def stencil(self, arg, var_accesses=None):
328  '''Not implemented.
329 
330  :param arg: the kernel argument with which the stencil is associated.
331  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
332  :param var_accesses: an unused optional argument that stores \
333  information about variable accesses.
334  :type var_accesses: :\
335  py:class:`psyclone.core.VariablesAccessInfo`
336 
337  :raises NotImplementedError: as this method is not implemented.
338 
339  '''
340  raise NotImplementedError("TODO #928: stencil not implemented")
341 
342  def operator(self, arg, var_accesses=None):
343  '''Create an LFRic operator argument and an ncells argument and add
344  them to the symbol table and argument list. Also declare the
345  associated 'fs_from', 'fs_to' symbols if they have not already
346  been declared so that they can be used to dimension the
347  operator symbol (as well as ncells).
348 
349  :param arg: the operator to add.
350  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
351  :param var_accesses: an unused optional argument that stores
352  information about variable accesses.
353  :type var_accesses: :py:class:`psyclone.core.VariablesAccessInfo`
354 
355  :raises NotImplementedError: if the datatype of the field is
356  not supported.
357 
358  '''
359  fs_from_name = arg.function_space_from.orig_name
360 
361  ndf_symbol_from = self._symtab_symtab.find_or_create_tag(
362  f"ndf_{fs_from_name}", fs=fs_from_name,
363  symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
364  interface=self._read_access_read_access)
365  fs_to_name = arg.function_space_to.orig_name
366  ndf_symbol_to = self._symtab_symtab.find_or_create_tag(
367  f"ndf_{fs_to_name}", fs=fs_to_name,
368  symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
369  interface=self._read_access_read_access)
370 
371  # We may already have a symbol for this argument as we add it for each
372  # operator (TODO #2074).
373  ncells = self._symtab_symtab.find_or_create(
374  "ncell_3d",
375  symbol_type=LFRicTypes("NumberOfCellsDataSymbol"),
376  interface=self._read_access_read_access)
377  self._arglist_arglist_arglist.append(ncells)
378 
379  op_arg_symbol = self._symtab_symtab.find_or_create_tag(
380  arg.name, symbol_type=LFRicTypes("OperatorDataSymbol"),
381  dims=[Reference(ndf_symbol_from), Reference(ndf_symbol_to),
382  Reference(ncells)],
383  fs_from=fs_from_name, fs_to=fs_to_name,
384  interface=ArgumentInterface(INTENT_MAPPING[arg.intent]))
385 
386  self._arg_index_to_metadata_index_arg_index_to_metadata_index[len(self._arglist_arglist_arglist)] = (
387  arg.metadata_index)
388  self._arglist_arglist_arglist.append(op_arg_symbol)
389 
390  def cma_operator(self, arg, var_accesses=None):
391  '''Not implemented.
392 
393  :param arg: the CMA operator argument.
394  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
395  :param var_accesses: an unused optional argument that stores \
396  information about variable accesses.
397  :type var_accesses: :\
398  py:class:`psyclone.core.VariablesAccessInfo`
399 
400  :raises NotImplementedError: as this method is not implemented.
401 
402  '''
403  raise NotImplementedError("TODO #928: cma_operator not implemented")
404 
405  def scalar(self, scalar_arg, var_accesses=None):
406  '''Create an LFRic scalar argument and add it to the symbol table and
407  argument list.
408 
409  :param scalar_arg: the scalar to add.
410  :type scalar_arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
411  :param var_accesses: an unused optional argument that stores \
412  information about variable accesses.
413  :type var_accesses: :\
414  py:class:`psyclone.core.VariablesAccessInfo`
415 
416  :raises NotImplementedError: if the datatype of the scalar is \
417  not supported.
418 
419  '''
420  mapping = {
421  "integer": LFRicTypes("LFRicIntegerScalarDataSymbol"),
422  "real": LFRicTypes("LFRicRealScalarDataSymbol"),
423  "logical": LFRicTypes("LFRicLogicalScalarDataSymbol")}
424  try:
425  symbol = self._symtab_symtab.find_or_create_tag(
426  scalar_arg.name,
427  symbol_type=mapping[scalar_arg.intrinsic_type],
428  interface=ArgumentInterface(INTENT_MAPPING[scalar_arg.intent]))
429  except KeyError as info:
430  message = (
431  f"scalar of type '{scalar_arg.intrinsic_type}' not implemented"
432  f" in KernelInterface class.")
433  raise NotImplementedError(message) from info
434  self._arg_index_to_metadata_index_arg_index_to_metadata_index[len(self._arglist_arglist_arglist)] = (
435  scalar_arg.metadata_index)
436  self._arglist_arglist_arglist.append(symbol)
437 
438  def fs_common(self, function_space, var_accesses=None):
439  '''Create any arguments that are common to a particular function
440  space. At this time the only common argument is the number of
441  degrees of freedom. Create the associated LFRic symbol, and
442  add it to the symbol table and argument list.
443 
444  :param function_space: the function space for any common arguments.
445  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
446  :param var_accesses: an unused optional argument that stores \
447  information about variable accesses.
448  :type var_accesses: :\
449  py:class:`psyclone.core.VariablesAccessInfo`
450 
451  '''
452  fs_name = function_space.orig_name
453  ndf_symbol = self._symtab_symtab.find_or_create_tag(
454  f"ndf_{fs_name}", fs=fs_name,
455  symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
456  interface=self._read_access_read_access)
457  self._arglist_arglist_arglist.append(ndf_symbol)
458 
459  def fs_intergrid(self, function_space, var_accesses=None):
460  '''Not implemented.
461 
462  :param arg: the CMA operator argument.
463  :type arg: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
464  :param var_accesses: an unused optional argument that stores \
465  information about variable accesses.
466  :type var_accesses: :\
467  py:class:`psyclone.core.VariablesAccessInfo`
468 
469  :raises NotImplementedError: as this method is not implemented.
470 
471  '''
472  raise NotImplementedError("TODO #928: fs_intergrid not implemented")
473 
474  def fs_compulsory_field(self, function_space, var_accesses=None):
475  '''Create any arguments that are compulsory for a field on a
476  particular function space. At this time the compulsory
477  arguments are the unique number of degrees of freedom and the
478  dofmap. Create the associated LFRic symbol, and add it to the
479  symbol table and argument list. Also declare the number of
480  degrees of freedom and add to the symbol table if one has not
481  yet been added.
482 
483  :param function_space: the function space for any compulsory arguments.
484  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
485  :param var_accesses: an unused optional argument that stores \
486  information about variable accesses.
487  :type var_accesses: :\
488  py:class:`psyclone.core.VariablesAccessInfo`
489 
490  '''
491  fs_name = function_space.orig_name
492  undf_symbol = self._symtab_symtab.find_or_create_tag(
493  f"undf_{fs_name}", fs=fs_name,
494  symbol_type=LFRicTypes("NumberOfUniqueDofsDataSymbol"),
495  interface=self._read_access_read_access)
496  self._arglist_arglist_arglist.append(undf_symbol)
497 
498  fs_name = function_space.orig_name
499  ndf_symbol = self._symtab_symtab.find_or_create_tag(
500  f"ndf_{fs_name}", fs=fs_name,
501  symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
502  interface=self._read_access_read_access)
503 
504  dofmap_symbol = self._symtab_symtab.find_or_create_tag(
505  f"dofmap_{fs_name}", fs=fs_name,
506  symbol_type=LFRicTypes("DofMapDataSymbol"),
507  dims=[Reference(ndf_symbol)], interface=self._read_access_read_access)
508  self._arglist_arglist_arglist.append(dofmap_symbol)
509 
510  def banded_dofmap(self, function_space, var_accesses=None):
511  '''Not implemented.
512 
513  :param function_space: the function space for this dofmap.
514  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
515  :param var_accesses: an unused optional argument that stores \
516  information about variable accesses.
517  :type var_accesses: :\
518  py:class:`psyclone.core.VariablesAccessInfo`
519 
520  :raises NotImplementedError: as this method is not implemented.
521 
522  '''
523  raise NotImplementedError("TODO #928: banded_dofmap not implemented")
524 
525  def indirection_dofmap(self, function_space, operator=None,
526  var_accesses=None):
527  '''Not implemented.
528 
529  :param function_space: the function space for this dofmap.
530  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
531  :param operator: the CMA operator.
532  :type operator: :py:class:`psyclone.dynamo0p3.DynKernelArgument`
533  :param var_accesses: an unused optional argument that stores \
534  information about variable accesses.
535  :type var_accesses: :\
536  py:class:`psyclone.core.VariablesAccessInfo`
537 
538  :raises NotImplementedError: as this method is not implemented.
539 
540  '''
541  raise NotImplementedError(
542  "TODO #928: indirection_dofmap not implemented")
543 
544  def basis(self, function_space, var_accesses=None):
545  '''Create an LFRic basis function argument and add it to the symbol
546  table and argument list.
547 
548  :param function_space: the function space for this basis function.
549  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
550  :param var_accesses: an unused optional argument that stores \
551  information about variable accesses.
552  :type var_accesses: :\
553  py:class:`psyclone.core.VariablesAccessInfo`
554 
555  '''
556  basis_name_func = function_space.get_basis_name
557  # This import must be placed here to avoid circular dependencies
558  # pylint: disable=import-outside-toplevel
559  from psyclone.dynamo0p3 import DynBasisFunctions
560  first_dim_value_func = DynBasisFunctions.basis_first_dim_value
561  self._create_basis_create_basis(function_space, self.basis_mappingbasis_mapping,
562  basis_name_func, first_dim_value_func)
563 
564  def diff_basis(self, function_space, var_accesses=None):
565  '''Create an LFRic differential basis function argument and add it to
566  the symbol table and argument list.
567 
568  :param function_space: the function space for this \
569  differential basis function.
570  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
571  :param var_accesses: an unused optional argument that stores \
572  information about variable accesses.
573  :type var_accesses: :\
574  py:class:`psyclone.core.VariablesAccessInfo`
575 
576  '''
577  basis_name_func = function_space.get_diff_basis_name
578  # This import must be placed here to avoid circular dependencies
579  # pylint: disable=import-outside-toplevel
580  from psyclone.dynamo0p3 import DynBasisFunctions
581  first_dim_value_func = DynBasisFunctions.diff_basis_first_dim_value
582  self._create_basis_create_basis(function_space, self.diff_basis_mappingdiff_basis_mapping,
583  basis_name_func, first_dim_value_func)
584 
585  def field_bcs_kernel(self, function_space, var_accesses=None):
586  '''Create the boundary-dofs mask argument required for the
587  enforce_bc_code kernel. Adds it to the symbol table and the argument
588  list.
589 
590  :param function_space: the function space for this boundary condition.
591  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
592  :param var_accesses: an unused optional argument that stores
593  information about variable accesses.
594  :type var_accesses: :py:class:`psyclone.core.VariablesAccessInfo`
595 
596  :raises InternalError: if the kernel does not have a single field as
597  argument.
598  :raises InternalError: if the field argument is not on the
599  'ANY_SPACE_1' function space.
600 
601  '''
602  # Sanity check - expect the enforce_bc_code to have a single argument
603  # that is a field.
604  if (len(self._kern_kern.arguments.args) != 1 or
605  not self._kern_kern.arguments.args[0].is_field):
606  const = LFRicConstants()
607  raise InternalError(
608  f"Kernel '{self._kern.name}' applies boundary conditions to a "
609  f"field and therefore should have a single, field argument "
610  f"(one of {const.VALID_FIELD_NAMES}) but got "
611  f"{[arg.argument_type for arg in self._kern.arguments.args]}")
612  farg = self._kern_kern.arguments.args[0]
613  fspace = farg.function_space
614  # Sanity check - expect the field argument to the enforce_bc_code
615  # kernel to be on the ANY_SPACE_1 space.
616  if fspace.orig_name != "any_space_1":
617  raise InternalError(
618  f"Kernel '{self._kern.name}' applies boundary conditions to a "
619  f"field but the supplied argument, '{farg.name}', is on "
620  f"'{fspace.orig_name}' rather than the expected 'ANY_SPACE_1'")
621 
622  ndf_symbol = self._symtab_symtab.find_or_create_tag(
623  f"ndf_{fspace.orig_name}", fs=fspace.orig_name,
624  symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
625  interface=self._read_access_read_access)
626 
627  # Add the boundary-dofs array argument.
628  sym = self._symtab_symtab.find_or_create_tag(
629  f"boundary_dofs_{farg.name}",
630  interface=self._read_access_read_access,
631  symbol_type=LFRicTypes("VerticalBoundaryDofMaskDataSymbol"),
632  dims=[Reference(ndf_symbol), 2])
633  self._arglist_arglist_arglist.append(sym)
634 
635  def operator_bcs_kernel(self, function_space, var_accesses=None):
636  '''Not implemented.
637 
638  :param function_space: the function space for this bcs kernel
639  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
640  :param var_accesses: an unused optional argument that stores \
641  information about variable accesses.
642  :type var_accesses: :\
643  py:class:`psyclone.core.VariablesAccessInfo`
644 
645  :raises NotImplementedError: as this method is not implemented.
646 
647  '''
648  raise NotImplementedError(
649  "TODO #928: operator_bcs_kernel not implemented")
650 
651  def ref_element_properties(self, var_accesses=None):
652  ''' Properties associated with the reference element
653 
654  :param var_accesses: an unused optional argument that stores \
655  information about variable accesses.
656  :type var_accesses: :\
657  py:class:`psyclone.core.VariablesAccessInfo`
658 
659  '''
660  # This callback does not contribute any kernel arguments
661 
662  def mesh_properties(self, var_accesses=None):
663  ''' Properties associated with the mesh
664 
665  :param var_accesses: an unused optional argument that stores \
666  information about variable accesses.
667  :type var_accesses: :\
668  py:class:`psyclone.core.VariablesAccessInfo`
669 
670  '''
671  # This callback does not contribute any kernel arguments
672 
673  def quad_rule(self, var_accesses=None):
674  '''Create LFRic arguments associated with the required quadrature, if
675  they do not already exist, and add them to the symbol table
676  and argument list. The arguments depend on the type of
677  quadrature requested.
678 
679  :param var_accesses: an unused optional argument that stores \
680  information about variable accesses.
681  :type var_accesses: :\
682  py:class:`psyclone.core.VariablesAccessInfo`
683 
684  :raises InternalError: if an unsupported quadrature shape is \
685  found.
686 
687  '''
688  # The kernel captures all the required quadrature shapes
689  for shape in self._kern.qr_rules:
690  if shape == "gh_quadrature_xyoz":
691  nqp_xy = self._symtab.find_or_create_tag(
692  "nqp_xy",
693  symbol_type=LFRicTypes("NumberOfQrPointsInXyDataSymbol"),
694  interface=self._read_access)
695  nqp_z = self._symtab.find_or_create_tag(
696  "nqp_z",
697  symbol_type=LFRicTypes("NumberOfQrPointsInZDataSymbol"),
698  interface=self._read_access)
699  weights_xy = self._symtab.find_or_create_tag(
700  "weights_xy",
701  symbol_type=LFRicTypes("QrWeightsInXyDataSymbol"),
702  dims=[Reference(nqp_xy)], interface=self._read_access)
703  weights_z = self._symtab.find_or_create_tag(
704  "weights_z",
705  symbol_type=LFRicTypes("QrWeightsInZDataSymbol"),
706  dims=[Reference(nqp_z)], interface=self._read_access)
707  self._arglist.extend([nqp_xy, nqp_z, weights_xy, weights_z])
708  elif shape == "gh_quadrature_face":
709  nfaces = self._symtab.find_or_create_tag(
710  "nfaces",
711  symbol_type=LFRicTypes("NumberOfFacesDataSymbol"),
712  interface=self._read_access)
713  nqp = self._symtab.find_or_create_tag(
714  "nqp_faces",
715  symbol_type=LFRicTypes(
716  "NumberOfQrPointsInFacesDataSymbol"),
717  interface=self._read_access)
718  weights = self._symtab.find_or_create_tag(
719  "weights_faces",
720  symbol_type=LFRicTypes("QrWeightsInFacesDataSymbol"),
721  dims=[Reference(nqp)], interface=self._read_access)
722  self._arglist.extend([nfaces, nqp, weights])
723  elif shape == "gh_quadrature_edge":
724  nedges = self._symtab.find_or_create_tag(
725  "nedges",
726  symbol_type=LFRicTypes("NumberOfEdgesDataSymbol"),
727  interface=self._read_access)
728  nqp = self._symtab.find_or_create_tag(
729  "nqp_edges",
730  symbol_type=LFRicTypes(
731  "NumberOfQrPointsInEdgesDataSymbol"),
732  interface=self._read_access)
733  weights = self._symtab.find_or_create_tag(
734  "weights_edges",
735  symbol_type=LFRicTypes("QrWeightsInEdgesDataSymbol"),
736  dims=[Reference(nqp)], interface=self._read_access)
737  self._arglist.extend([nedges, nqp, weights])
738  else:
739  raise InternalError(f"Unsupported quadrature shape '{shape}' "
740  f"found in kernel_interface.")
741 
742  def _create_basis(self, function_space, mapping, basis_name_func,
743  first_dim_value_func):
744  '''Internal utility to create an LFRic basis or differential basis
745  function argument specific to the particular quadrature that
746  is being used and add it to the symbol table and argument
747  list. Also declare the associated "ndf" symbol and any
748  quadrature-specific symbols if they have not already been
749  declared so that they can be used to dimension the basis or
750  differential basis symbol.
751 
752  This utility function is used to avoid code replication as the
753  structure of a basis function is very similar to the structure
754  of a differential basis function.
755 
756  :param function_space: the function space that this basis or \
757  differential basis function is on.
758  :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
759  :param dict mapping: a mapping from quadrature type to basis \
760  or differential basis class name.
761  :param method basis_name_func: a method that returns the name \
762  of the basis or differential basis function for the \
763  current function space.
764  :param function first_dim_value_func: a function that returns \
765  the size of the first dimension of the basis or \
766  differential basis function for the current function \
767  space.
768 
769  :raises NotImplementedError: if an evaluator shape is found \
770  that is not a quadrature shape (currently just \
771  'gh_evaluator').
772  :raises InternalError: if the supplied evaluator shape is not \
773  recognised.
774 
775  '''
776  # pylint: disable=too-many-locals
777  const = LFRicConstants()
778  for shape in self._kern.eval_shapes:
779  fs_name = function_space.orig_name
780  ndf_symbol = self._symtab.find_or_create_tag(
781  f"ndf_{fs_name}",
782  symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
783  fs=fs_name, interface=self._read_access)
784 
785  # Create the qr tag by appending the last part of the shape
786  # name to "qr_".
787  quad_name = shape.split("_")[-1]
788  basis_tag = basis_name_func(qr_var="qr_"+quad_name)
789  if shape == "gh_quadrature_xyoz":
790  nqp_xy = self._symtab.find_or_create_tag(
791  "nqp_xy",
792  symbol_type=LFRicTypes("NumberOfQrPointsInXyDataSymbol"),
793  interface=self._read_access)
794  nqp_z = self._symtab.find_or_create_tag(
795  "nqp_z",
796  symbol_type=LFRicTypes("NumberOfQrPointsInZDataSymbol"),
797  interface=self._read_access)
798  type_name = mapping["gh_quadrature_xyoz"]
799  arg = LFRicTypes(type_name)(
800  basis_tag, [int(first_dim_value_func(function_space)),
801  Reference(ndf_symbol), Reference(nqp_xy),
802  Reference(nqp_z)],
803  fs_name, interface=self._read_access)
804  elif shape == "gh_quadrature_face":
805  nfaces = self._symtab.find_or_create_tag(
806  "nfaces",
807  symbol_type=LFRicTypes("NumberOfFacesDataSymbol"),
808  interface=self._read_access)
809  nqp = self._symtab.find_or_create_tag(
810  "nqp_faces",
811  symbol_type=LFRicTypes(
812  "NumberOfQrPointsInFacesDataSymbol"),
813  interface=self._read_access)
814  type_name = mapping["gh_quadrature_face"]
815  arg = LFRicTypes(type_name)(
816  basis_tag, [int(first_dim_value_func(function_space)),
817  Reference(ndf_symbol), Reference(nqp),
818  Reference(nfaces)],
819  fs_name, interface=self._read_access)
820  elif shape == "gh_quadrature_edge":
821  nedges = self._symtab.find_or_create_tag(
822  "nedges",
823  symbol_type=LFRicTypes("NumberOfEdgesDataSymbol"),
824  interface=self._read_access)
825  nqp = self._symtab.find_or_create_tag(
826  "nqp_edges",
827  symbol_type=LFRicTypes(
828  "NumberOfQrPointsInEdgesDataSymbol"),
829  interface=self._read_access)
830  type_name = mapping["gh_quadrature_edge"]
831  arg = LFRicTypes(type_name)(
832  basis_tag, [int(first_dim_value_func(function_space)),
833  Reference(ndf_symbol), Reference(nqp),
834  Reference(nedges)],
835  fs_name, interface=self._read_access)
836  elif shape in const.VALID_EVALUATOR_SHAPES:
837  # Need a (diff) basis array for each target space upon
838  # which the basis functions have been
839  # evaluated. _kern.eval_targets is a dict where the
840  # values are 2-tuples of (FunctionSpace, argument).
841  for _, _ in self._kern.eval_targets.items():
842  raise NotImplementedError(
843  "TODO #928: Evaluator shapes not implemented in "
844  "kernel_interface class.")
845  else:
846  raise InternalError(
847  f"Unrecognised quadrature or evaluator shape '{shape}'. "
848  f"Expected one of: {const.VALID_EVALUATOR_SHAPES}.")
849  self._symtab.add(arg)
850  self._arglist.append(arg)
def append(self, var_name, var_accesses=None, var_access_name=None, mode=AccessType.READ, metadata_posn=None)
def extend(self, list_var_name, var_accesses=None, mode=AccessType.READ, list_metadata_posn=None)
def banded_dofmap(self, function_space, var_accesses=None)
def field_vector(self, argvect, var_accesses=None)
def _create_basis(self, function_space, mapping, basis_name_func, first_dim_value_func)
def fs_common(self, function_space, var_accesses=None)
def basis(self, function_space, var_accesses=None)
def indirection_dofmap(self, function_space, operator=None, var_accesses=None)
def diff_basis(self, function_space, var_accesses=None)
def fs_intergrid(self, function_space, var_accesses=None)
def stencil_unknown_extent(self, arg, var_accesses=None)
def scalar(self, scalar_arg, var_accesses=None)
def stencil_unknown_direction(self, arg, var_accesses=None)
def field_bcs_kernel(self, function_space, var_accesses=None)
def operator_bcs_kernel(self, function_space, var_accesses=None)
def fs_compulsory_field(self, function_space, var_accesses=None)