Add temporary arrays transformation into pointers, and field_new/field_delete

ecossevin · ecossevin · commit cacb0238ecfa · 2024-04-25T14:02:05.000+02:00
diff --git a/transformations/tests/test_parallel_routine_dispatch.py b/transformations/tests/test_parallel_routine_dispatch.py
@@ -11,7 +11,7 @@
 import pytest
 
 from loki.frontend import available_frontends, OMNI
-from loki import Sourcefile, FindNodes, CallStatement
+from loki import Sourcefile, FindNodes, CallStatement, fgen, Conditional
 
 from transformations.parallel_routine_dispatch import ParallelRoutineDispatchTransformation
 
@@ -33,8 +33,71 @@ def test_parallel_routine_dispatch_dr_hook(here, frontend):
     transformation = ParallelRoutineDispatchTransformation()
     transformation.apply(source['dispatch_routine'])
 
-    calls = FindNodes(CallStatement).visit(routine.body)
+    calls = [call for call in FindNodes(CallStatement).visit(routine.body) if call.name.name=='DR_HOOK']
+    assert len(calls) == 8
+
+@pytest.mark.parametrize('frontend', available_frontends(skip=[OMNI]))
+def test_parallel_routine_dispatch_decl_local_arrays(here, frontend):
+
+    source = Sourcefile.from_file(here/'sources/projParallelRoutineDispatch/dispatch_routine.F90', frontend=frontend)
+    routine = source['dispatch_routine']
+
+    transformation = ParallelRoutineDispatchTransformation()
+    transformation.apply(source['dispatch_routine'])
+    var_lst=["YL_ZRDG_CVGQ", "ZRDG_CVGQ", "YL_ZRDG_MU0LU", "ZRDG_MU0LU", "YL_ZRDG_MU0M", "ZRDG_MU0M", "YL_ZRDG_MU0N", "ZRDG_MU0N", "YL_ZRDG_MU0", "ZRDG_MU0"]
+    dcls = [dcl for dcl in routine.declarations if dcl.symbols[0].name in var_lst]
+    str_dcls = ""
+    for dcl in dcls:
+        str_dcls += fgen(dcl)+"\n"
+    assert str_dcls == """CLASS(FIELD_3RB), POINTER :: YL_ZRDG_CVGQ => NULL()
+REAL(KIND=JPRB), POINTER :: ZRDG_CVGQ(:, :, :) => NULL()
+CLASS(FIELD_2RB), POINTER :: YL_ZRDG_MU0LU => NULL()
+REAL(KIND=JPRB), POINTER :: ZRDG_MU0LU(:, :) => NULL()
+CLASS(FIELD_2RB), POINTER :: YL_ZRDG_MU0M => NULL()
+REAL(KIND=JPRB), POINTER :: ZRDG_MU0M(:, :) => NULL()
+CLASS(FIELD_2RB), POINTER :: YL_ZRDG_MU0N => NULL()
+REAL(KIND=JPRB), POINTER :: ZRDG_MU0N(:, :) => NULL()
+CLASS(FIELD_2RB), POINTER :: YL_ZRDG_MU0 => NULL()
+REAL(KIND=JPRB), POINTER :: ZRDG_MU0(:, :) => NULL()
+"""
+
+@pytest.mark.parametrize('frontend', available_frontends(skip=[OMNI]))
+def test_parallel_routine_dispatch_decl_field_create_delete(here, frontend):
+
+    source = Sourcefile.from_file(here/'sources/projParallelRoutineDispatch/dispatch_routine.F90', frontend=frontend)
+    routine = source['dispatch_routine']
+
+    transformation = ParallelRoutineDispatchTransformation()
+    transformation.apply(source['dispatch_routine'])
+
+    var_lst = ["YL_ZRDG_CVGQ", "ZRDG_CVGQ", "YL_ZRDG_MU0LU", "ZRDG_MU0LU", "YL_ZRDG_MU0M", "ZRDG_MU0M", "YL_ZRDG_MU0N", "ZRDG_MU0N", "YL_ZRDG_MU0", "ZRDG_MU0"]
+    field_create = ["CALL FIELD_NEW(YL_ZRDG_CVGQ, UBOUNDS=(/ YDCPG_OPTS%KLON, YDCPG_OPTS%KFLEVG, YDCPG_OPTS%KGPBLKS /), LBOUNDS=(/ 0, 1 /),  &\n& PERSISTENT=.true.)",
+                "CALL FIELD_NEW(YL_ZRDG_MU0N, UBOUNDS=(/ YDCPG_OPTS%KLON, YDCPG_OPTS%KGPBLKS /), PERSISTENT=.true.)",
+                "CALL FIELD_NEW(YL_ZRDG_MU0LU, UBOUNDS=(/ YDCPG_OPTS%KLON, YDCPG_OPTS%KGPBLKS /), PERSISTENT=.true.)",
+                "CALL FIELD_NEW(YL_ZRDG_MU0, UBOUNDS=(/ YDCPG_OPTS%KLON, YDCPG_OPTS%KGPBLKS /), PERSISTENT=.true.)",
+                "CALL FIELD_NEW(YL_ZRDG_MU0M, UBOUNDS=(/ YDCPG_OPTS%KLON, YDCPG_OPTS%KGPBLKS /), PERSISTENT=.true.)"
+                ]
+
+    calls = [call for call in FindNodes(CallStatement).visit(routine.body) if call.name.name=="FIELD_NEW"]
     assert len(calls) == 5
-    assert [str(call.name).lower() for call in calls] == [
-        'dr_hook', 'dr_hook', 'cpphinp', 'dr_hook', 'dr_hook'
-    ]
+    for call in calls:
+        assert fgen(call) in field_create
+    
+    field_delete = ["IF (ASSOCIATED(YL_ZRDG_CVGQ)) CALL FIELD_DELETE(YL_ZRDG_CVGQ)",
+                "IF (ASSOCIATED(YL_ZRDG_MU0LU)) CALL FIELD_DELETE(YL_ZRDG_MU0LU)",
+                "IF (ASSOCIATED(YL_ZRDG_MU0M)) CALL FIELD_DELETE(YL_ZRDG_MU0M)",
+                "IF (ASSOCIATED(YL_ZRDG_MU0)) CALL FIELD_DELETE(YL_ZRDG_MU0)",
+                "IF (ASSOCIATED(YL_ZRDG_MU0N)) CALL FIELD_DELETE(YL_ZRDG_MU0N)"
+                ]
+
+    conds = [cond for cond in FindNodes(Conditional).visit(routine.body)]
+    conditional = []
+    for cond in conds:
+        for call in FindNodes(CallStatement).visit(cond):
+                if call.name.name=="FIELD_DELETE":
+                    conditional.append(cond)
+
+    assert len(conditional) == 5
+    for cond in conditional:
+        assert fgen(cond) in field_delete
+    breakpoint()
diff --git a/transformations/transformations/parallel_routine_dispatch.py b/transformations/transformations/parallel_routine_dispatch.py
@@ -11,17 +11,36 @@
     FindNodes, nodes as ir
 )
 from loki.transform import Transformation
+from loki import (
+    FindVariables, DerivedType, SymbolAttributes,
+    Array, single_variable_declaration, Transformer
+)
 
 __all__ = ['ParallelRoutineDispatchTransformation']
 
 
 class ParallelRoutineDispatchTransformation(Transformation):
 
+    def __init__(self):
+        self.horizontal = [
+            "KLON", "YDCPG_OPTS%KLON", "YDGEOMETRY%YRDIM%NPROMA",
+            "KPROMA", "YDDIM%NPROMA", "NPROMA"
+    ]
+        # CALL FIELD_NEW (YL_ZA, UBOUNDS=[KLON, KFLEVG, KGPBLKS], LBOUNDS=[1, 0, 1], PERSISTENT=.TRUE.)
+        self.new_calls = []
+        # IF (ASSOCIATED (YL_ZA)) CALL FIELD_DELETE (YL_ZA)
+        self.delete_calls = []
+        self.routine_map_temp = {} 
+
     def transform_subroutine(self, routine, **kwargs):
         with pragma_regions_attached(routine):
             for region in FindNodes(ir.PragmaRegion).visit(routine.body):
                 if is_loki_pragma(region.pragma):
                     self.process_parallel_region(routine, region)
+        single_variable_declaration(routine)
+        self.add_temp(routine)
+        self.add_field(routine)
+        #call add_arrays etc...
 
     def process_parallel_region(self, routine, region):
         pragma_content = region.pragma.content.split(maxsplit=1)
@@ -41,6 +60,14 @@ def process_parallel_region(self, routine, region):
         region.prepend(dr_hook_calls[0])
         region.append(dr_hook_calls[1])
 
+        region_map_temp= self.decl_local_array(routine, region)
+        
+        for var_name in region_map_temp:
+            if var_name not in self.routine_map_temp:
+                self.routine_map_temp[var_name]=region_map_temp[var_name]
+
+
+
     @staticmethod
     def create_dr_hook_calls(scope, cdname, handle):
         dr_hook_calls = []
@@ -56,3 +83,99 @@ def create_dr_hook_calls(scope, cdname, handle):
                 )
             ]
         return dr_hook_calls
+
+#    @staticmethod
+#    def get_local_array(routine):
+
+    def create_field_new_delete(self, routine, var, field_ptr_var):
+        # Create the FIELD_NEW call
+        var_shape = var.shape
+        ubounds = [d.upper if isinstance(d, sym.RangeIndex) else d for d in var_shape]
+        ubounds += [sym.Variable(name="KGPBLKS", parent=routine.variable_map["YDCPG_OPTS"])]
+        # NB: This is presumably what the condition should look like, however, the
+        #     logic is flawed in to_parallel and it will only insert lbounds if _the last_
+        #     dimension has an lbound. We emulate this with the second line here to
+        #     generate identical results, but this line should probably not be there
+        has_lbounds = any(isinstance(d, sym.RangeIndex) for d in var_shape)
+        has_lbounds = has_lbounds and isinstance(var_shape[-1], sym.RangeIndex)
+        if has_lbounds:
+            lbounds = [
+                d.lower if isinstance(d, sym.RangeIndex) else sym.IntLiteral(0)
+                for d in var_shape
+            ]
+            kwarguments = (
+                ('UBOUNDS', sym.LiteralList(ubounds)),
+                ('LBOUNDS', sym.LiteralList(lbounds)),
+                ('PERSISTENT', sym.LogicLiteral(True))
+            )
+        else:
+            kwarguments = (
+                ('UBOUNDS', sym.LiteralList(ubounds)),
+                ('PERSISTENT', sym.LogicLiteral(True))
+            )
+        self.new_calls += [ir.CallStatement(
+            name=sym.Variable(name='FIELD_NEW', scope=routine),
+            arguments=(field_ptr_var,),
+            kwarguments=kwarguments
+        )]
+
+        # Create the FIELD_DELETE CALL
+        call = ir.CallStatement(sym.Variable(name='FIELD_DELETE', scope=routine), arguments=(field_ptr_var,))
+        condition = sym.InlineCall(sym.Variable(name='ASSOCIATED'), parameters=(field_ptr_var,))
+        self.delete_calls += [ir.Conditional(condition=condition, inline=True, body=(call,))]
+
+    def decl_local_array(self, routine, region):
+        temp_arrays = [var for var in FindVariables(Array).visit(region)  if isinstance(var, Array) and not var.name_parts[0] in routine.arguments and var.shape[0] in self.horizontal]
+        region_map_temp={}        
+        #check if first dim NPROMA ?
+        for var in temp_arrays:
+            var_type = var.type
+            var_shape = var.shape
+
+            dim = len(var_shape) + 1 # Temporary dimensions + block
+
+            # The FIELD_{d}RB variable
+            field_ptr_type = SymbolAttributes(
+                dtype=DerivedType(f'FIELD_{dim}RB'),
+                pointer=True, polymorphic=True, initial="NULL()"
+            )
+            field_ptr_var = sym.Variable(name=f'YL_{var.name}', type=field_ptr_type, scope=routine)
+
+            # Create a pointer instead of the array
+            shape = (sym.RangeIndex((None, None)),) * dim
+          #  var.type = var_type.clone(pointer=True, shape=shape)
+            local_ptr_var = var.clone(dimensions=shape)
+
+            region_map_temp[var.name]=[field_ptr_var,local_ptr_var]
+            self.create_field_new_delete(routine, var, field_ptr_var)
+        return(region_map_temp)
+
+    def add_temp(self, routine):
+        # Replace temporary declaration by pointer to array and pointer to field_api object
+        map_dcl = {}
+        for decl in routine.declarations:
+            if len(decl.symbols) == 1:
+                var = decl.symbols[0]
+                if var.name in self.routine_map_temp:
+                    new_vars = self.routine_map_temp[var.name]
+                    new_vars[1].type = new_vars[1].type.clone(pointer=True, shape=new_vars[1].dimensions, initial="NULL()")
+                    map_dcl.update({decl : (ir.VariableDeclaration(symbols=(new_vars[0],)), ir.VariableDeclaration(symbols=(new_vars[1],)))})
+            else:
+                raise Exception("Declaration should have only one symbol, please run single_variable_declaration before calling add_temp function.")
+        routine.spec = Transformer(map_dcl).visit(routine.spec)
+        #decls_to_replace = [var for var in decl.var for decl in routine.declarations if var in self.routine_map_temp]
+    
+    def add_field(self, routine):
+        # Insert the field generation wrapped into a DR_HOOK call
+        dr_hook_calls = self.create_dr_hook_calls(
+            routine, cdname='CREATE_TEMPORARIES',
+            handle=sym.Variable(name='ZHOOK_HANDLE_FIELD_API', scope=routine)
+        )
+        routine.body.insert(2, (dr_hook_calls[0], ir.Comment(text=''), *self.new_calls, dr_hook_calls[1]))
+
+        # Insert the field deletion wrapped into a DR_HOOK call
+        dr_hook_calls = self.create_dr_hook_calls(
+            routine, cdname='DELETE_TEMPORARIES',
+            handle=sym.Variable(name='ZHOOK_HANDLE_FIELD_API', scope=routine)
+        )
+        routine.body.insert(-2,(dr_hook_calls[0], ir.Comment(text=''), *self.delete_calls, dr_hook_calls[1]))
