final clean-up, add documentation, etc

loki/backend/cgen.py

@@ -22,6 +22,9 @@
 
 
 class IntrinsicTypeC:
+    """
+    Mapping Fortran type to corresponding C type.
+    """
     # pylint: disable=abstract-method, unused-argument
 
     def __init__(self, *args, **kwargs):
@@ -45,6 +48,10 @@ def c_intrinsic_type(self, _type, *args, **kwargs):
 c_intrinsic_type = IntrinsicTypeC()
 
 class CCodeMapper(LokiStringifyMapper):
+    """
+    A :class:`StringifyMapper`-derived visitor for Pymbolic expression trees that converts an
+    expression to a string adhering to standardized C.
+    """
     # pylint: disable=abstract-method, unused-argument
 
     def __init__(self, c_intrinsic_type, *args, **kwargs):
@@ -103,7 +110,6 @@ def map_array_subscript(self, expr, enclosing_prec, *args, **kwargs):
                 if d:
                     index_str += self.format('[%s]', d)
             return self.format('%s%s', name_str, index_str)
-        # else:
         return self.format('%s', name_str)
 
     map_string_subscript = map_array_subscript
@@ -170,6 +176,9 @@ def visit_Module(self, o, **kwargs):
         return self.join_lines(spec, routines)
 
     def _subroutine_header(self, o, **kwargs):
+        """
+        Helper function/header for :func:`~loki.backend.CCodegen.visit_Subroutine`.
+        """
         # Some boilerplate imports...
         header = [self.format_line('#include <', name, '>') for name in self.standard_imports]
         # ...and imports from the spec
@@ -178,6 +187,9 @@ def _subroutine_header(self, o, **kwargs):
         return header
 
     def _subroutine_arguments(self, o, **kwargs):
+        """
+        Helper function/routine arguments for :func:`~loki.backend.CCodegen.visit_Subroutine`.
+        """
         var_keywords = []
         pass_by = []
         for a in o.arguments:
@@ -193,6 +205,9 @@ def _subroutine_arguments(self, o, **kwargs):
         return pass_by, var_keywords
 
     def _subroutine_declaration(self, o, **kwargs):
+        """
+        Helper function/function declaration part for :func:`~loki.backend.CCodegen.visit_Subroutine`.
+        """
         pass_by, var_keywords = self._subroutine_arguments(o, **kwargs)
         arguments = [f'{k}{self.visit(a.type, **kwargs)} {p}{a.name}'
                      for a, p, k in zip(o.arguments, pass_by, var_keywords)]
@@ -207,6 +222,9 @@ def _subroutine_declaration(self, o, **kwargs):
         return declaration
 
     def _subroutine_body(self, o, **kwargs):
+        """
+        Helper function/body for :func:`~loki.backend.CCodegen.visit_Subroutine`.
+        """
         self.depth += 1
 
         # body = ['{']
@@ -226,6 +244,9 @@ def _subroutine_body(self, o, **kwargs):
         return body
 
     def _subroutine_footer(self, o, **kwargs):
+        """
+        Helper function/footer for :func:`~loki.backend.CCodegen.visit_Subroutine`.
+        """
         footer = [self.format_line('}')]
         return footer
 
@@ -234,7 +255,7 @@ def visit_Subroutine(self, o, **kwargs):
         Format as:
 
           ...imports...
-          int <name>(<args>) {
+          <return_type> <name>(<args>) {
             ...spec without imports and argument declarations...
             ...body...
           }
@@ -427,6 +448,12 @@ def visit_SymbolAttributes(self, o, **kwargs):  # pylint: disable=unused-argumen
         return self.symgen.c_intrinsic_type(o)
 
     def visit_TypeDef(self, o, **kwargs):
+        """
+        Format type definition/struct as
+          struct <name> {
+            ...declarations...
+          };
+        """
         header = self.format_line('struct ', o.name.lower(), ' {')
         footer = self.format_line('};')
         self.depth += 1

loki/backend/cppgen.py

@@ -13,6 +13,9 @@
 
 
 class IntrinsicTypeCpp(IntrinsicTypeC):
+    """
+    Mapping Fortran type to corresponding C++ type.
+    """
 
     def c_intrinsic_type(self, _type, *args, **kwargs):
         if _type.dtype == BasicType.INTEGER:
@@ -25,6 +28,10 @@ def c_intrinsic_type(self, _type, *args, **kwargs):
 
 
 class CppCodeMapper(CCodeMapper):
+    """
+    A :class:`StringifyMapper`-derived visitor for Pymbolic expression trees that converts an
+    expression to a string adhering to standardized C++.
+    """
     # pylint: disable=abstract-method, unused-argument
     pass
 
@@ -47,9 +54,6 @@ def _subroutine_header(self, o, **kwargs):
         return header
 
     def _subroutine_arguments(self, o, **kwargs):
-        # opt_extern = kwargs.get('extern', False)
-        # if opt_extern:
-        #     return super()._subroutine_arguments(o, **kwargs)
         var_keywords = []
         pass_by = []
         for a in o.arguments:

loki/backend/cudagen.py

@@ -4,24 +4,28 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-# from loki.expression import Array
 from loki.types import DerivedType
 from loki.backend.cppgen import CppCodegen, CppCodeMapper, IntrinsicTypeCpp
-
-# from loki.tools import as_tuple
 from loki.ir import Import, FindNodes
 from loki.expression import Array
 
 __all__ = ['cudagen', 'CudaCodegen', 'CudaCodeMapper']
 
 
 class IntrinsicTypeCuda(IntrinsicTypeCpp):
+    """
+    Mapping Fortran type to corresponding CUDA type.
+    """
     pass
 
 cuda_intrinsic_type = IntrinsicTypeCuda()
 
 
 class CudaCodeMapper(CppCodeMapper):
+    """
+    A :class:`StringifyMapper`-derived visitor for Pymbolic expression trees that converts an
+    expression to a string adhering to standardized CUDA.
+    """
     # pylint: disable=abstract-method, unused-argument
     pass
 
@@ -84,8 +88,6 @@ def _subroutine_declaration(self, o, **kwargs):
                      for a, p, k in zip(o.arguments, pass_by, var_keywords)]
         opt_header = kwargs.get('header', False)
         end = ' {' if not opt_header else ';'
-        # check whether to return something and define function return type accordingly
-        ##
         prefix = ''
         if o.prefix and "global" in o.prefix[0].lower():
             prefix = '__global__ '
@@ -100,22 +102,16 @@ def _subroutine_declaration(self, o, **kwargs):
 
     def _subroutine_body(self, o, **kwargs):
         self.depth += 1
-
-        # body = ['{']
         # ...and generate the spec without imports and argument declarations
         body = [self.visit(o.spec, skip_imports=True, skip_argument_declarations=True, **kwargs)]
-
         # Fill the body
         body += [self.visit(o.body, **kwargs)]
-
         opt_extern = kwargs.get('extern', False)
         if opt_extern:
             body += [self.format_line('cudaDeviceSynchronize();')]
-
         # if something to be returned, add 'return <var>' statement
         if o.result_name is not None:
             body += [self.format_line(f'return {o.result_name.lower()};')]
-
         # Close everything off
         self.depth -= 1
         # footer = [self.format_line('}')]

loki/transformations/block_index_transformations.py

@@ -422,7 +422,10 @@ def process_kernel(self, routine, targets, exclude_arrays):
 
 class LowerBlockIndexTransformation(Transformation):
     """
-    ...
+    Transformation to lower the block index via appending the block index 
+    to variable dimensions/shape. However, this only handles variable
+    declarations/definitions. Therefore this transformation must always be followed by
+    the :any:`InjectBlockIndexTransformation`.
 
     Parameters
     ----------
@@ -505,29 +508,78 @@ def process(self, routine, targets, role):
             processed_routines += (call.routine.name,)
 
 
-
 class LowerBlockLoopTransformation(Transformation):
     """
-    ...
+    Lower the block loop to calls within this loop. 
+
+    For example, the following code:
+
+    .. code-block:: fortran
+
+        subroutine driver(nblks, ...)
+            ...
+            integer, intent(in) :: nblks
+            integer :: ibl
+            real :: var(jlon,nlev,nblks)
+
+            do ibl=1,nblks
+                call kernel2(var,...,nblks,ibl)
+            enddo
+            ...
+        end subroutine driver
+
+        subroutine kernel(var, ..., nblks, ibl)
+            ...
+            real :: var(jlon,nlev,nblks)
+
+            do jl=1,...
+                do jk=1,...
+                    var(jk,jl,ibl) = ...
+                end do
+            end do
+        end subroutine kernel
+
+    is transformed to:
+    
+    .. code-block:: fortran
+
+        subroutine driver(nblks, ...)
+            ...
+            integer, intent(in) :: nblks
+            integer :: ibl
+            real :: var(jlon,nlev,nblks)
+
+            call kernel2(var,..., nblks)
+            ...
+        end subroutine driver
+
+        subroutine kernel(var, ..., nblks)
+            ...
+            integer :: ibl
+            real :: var(jlon,nlev,nblks)
+            
+            do ibl=1,nblks
+                do jl=1,...
+                    do jk=1,...
+                        var(jk,jl,ibl) = ...
+                    end do
+                end do
+            end do
+        end subroutine kernel
 
     Parameters
     ----------
     block_dim : :any:`Dimension`
         :any:`Dimension` object describing the variable conventions used in code
         to define the blocking data dimension and iteration space.
-    recurse_to_kernels : bool, optional
-        Recurse/continue with/to (nested) kernels and lower the block index for those
-        as well (default: `False`).
     """
     # This trafo only operates on procedures
     item_filter = (ProcedureItem,)
 
-    def __init__(self, block_dim): # , recurse_to_kernels=False): # , key=None):
+    def __init__(self, block_dim):
         self.block_dim = block_dim
-        # self.remove_loop = True
 
     def transform_subroutine(self, routine, **kwargs):
-
         role = kwargs['role']
         targets = tuple(str(t).lower() for t in as_tuple(kwargs.get('targets', None)))
         if role == 'driver':
@@ -541,15 +593,10 @@ def arg_to_local_var(routine, var):
             type=routine.variable_map[var.name].type.clone(intent=None)),)
 
     def local_var(self, call, var):
-        # if var.name in call.arg_map:
-        #     print(f"arg to local var [1] {var} | {call.arg_map[var.name]}")
-        #     self.arg_to_local_var(call.routine, call.arg_map[var.name])
-        # elif var.name in call.routine.arguments:
         if var.name in call.routine.arguments:
             self.arg_to_local_var(call.routine, var)
         else:
             call.routine.variables += (var.clone(scope=call.routine),)
-            # (routine.variable_map[var.name].clone(scope=call.routine))
 
     @staticmethod
     def remove_openmp_pragmas(routine):
@@ -565,21 +612,16 @@ def generate_pragma(loop):
 
     def process_driver(self, routine, targets):
         # find block loops
-        # with pragmas_attached(routine, (ir.Loop, ir.CallStatement)):
-        # if True:
         loops = FindNodes(ir.Loop).visit(routine.body)
         loops = [loop for loop in loops if loop.variable == self.block_dim.index
                 or loop.variable in self.block_dim._index_aliases]
-        # if True:
-        # loop_map = {}
-        # ignore_routine = []
         driver_loop_map = {}
         to_local_var = {}
         processed_routines = ()
         calls = ()
         for loop in loops:
             lower_loop = False
-            for call in  FindNodes(ir.CallStatement).visit(loop.body): #visit(routine.body):
+            for call in  FindNodes(ir.CallStatement).visit(loop.body):
                 if str(call.name).lower() not in targets:
                     continue
                 lower_loop = True
@@ -590,7 +632,7 @@ def process_driver(self, routine, targets):
                 # replace/substitute variables according to the caller-callee argument map
                 loop_to_lower = SubstituteExpressions(call_arg_map).visit(loop_to_lower)
                 # remove calls that are not within targets # TODO: rather a hack to remove
-                #  "CALL TIMER%THREAD_LOG(TID, IGPC=ICEND)"
+                #  e.g., "CALL TIMER%THREAD_LOG(TID, IGPC=ICEND)"
                 calls_within_loop = [_call for _call in  FindNodes(ir.CallStatement).visit(loop_to_lower.body)
                         if str(_call.name).lower() not in targets]
                 loop_to_lower = Transformer({call: None for call in calls_within_loop}).visit(loop_to_lower)
@@ -630,10 +672,9 @@ def process_driver(self, routine, targets):
             if lower_loop:
                 driver_loop_map[loop] = loop.body
         routine.body = Transformer(driver_loop_map).visit(routine.body)
-        for call in calls: #  FindNodes(ir.CallStatement).visit(routine.body):
+        for call in calls:
             if str(call.name).lower() not in targets:
                 continue
-            # self.local_var(routine, call, loop.variable)
             for var in to_local_var[call.routine.name]:
                 self.local_var(call, var)
         # TODO: remove

loki/transformations/data_offload.py

@@ -833,8 +833,8 @@ def _get_symbols(self, successors):
         """
         Get module variables/symbols (grouped by routine/successor).
         """
-        defines_symbols = CaseInsensitiveDict() # {} # CaseInsensitiveDict()
-        uses_symbols = CaseInsensitiveDict() # {} # CaseInsensitiveDict()
+        defines_symbols = CaseInsensitiveDict()
+        uses_symbols = CaseInsensitiveDict()
         for item in successors:
             if not isinstance(item, ProcedureItem):
                 continue