(closes #2951) Consider dependencies in ArrayAssignment2LoopTrans

changelog

@@ -1,3 +1,6 @@
+   48) PR #3296 for #2951. Improves ArrayAssignment2LoopTrans to allow
+   for dependencies between LHS and RHS.
+
    47) PR #3332 for #3143 - fixes bug in DependenceUseChain so that
    previous access of a call argument does not return the call.
 

examples/nemo/eg4/sir_trans.py

@@ -52,5 +52,8 @@ def trans(psyir):
     sir_writer = SIRWriter()
     for subroutine in psyir.walk(Routine):
         print(f"Transforming subroutine: {subroutine.name}")
-        sir_code = sir_writer(subroutine)
-        print(sir_code)
+        try:
+            sir_code = sir_writer(subroutine)
+            print(sir_code)
+        except Exception as e:
+            print(f"Failed to transform {subroutine.name}: {e}")

examples/nemo/eg4/sir_trans_all.py

@@ -114,5 +114,8 @@ def trans(psyir):
 
         kern = fortran_writer(subroutine)
         print(kern)
-        kern = sir_writer(subroutine)
-        print(kern)
+        try:
+            kern = sir_writer(subroutine)
+            print(kern)
+        except Exception as e:
+            print(f"Failed to transform {subroutine.name}: {e}")

examples/nemo/eg4/sir_trans_loop.py

@@ -86,7 +86,10 @@ def trans(psyir):
                 if isinstance(child, Assignment):
                     hoist_trans.apply(child)
 
-        kern = sir_writer(subroutine)
-        # TODO issue #1854. There should be backend support for
-        # writing out SIR.
-        print(kern)
+        try:
+            kern = sir_writer(subroutine)
+            # TODO issue #1854. There should be backend support for
+            # writing out SIR.
+            print(kern)
+        except Exception as e:
+            print(f"Failed to transform {subroutine.name}: {e}")

examples/nemo/scripts/omp_gpu_trans.py

@@ -84,7 +84,6 @@
 
 NEMOV4_EXCLUSIONS = [
     "dynspg_ts.f90",
-    "tranxt.f90",
 ]
 
 SKIP_FOR_PERFORMANCE = [
@@ -191,12 +190,14 @@ def trans(psyir):
                 subroutine.name == 'dom_zgr' or
                 subroutine.name == 'dom_ngb'):
             continue
-        if subroutine.name == "solfrac_mod.f90":
+        if subroutine.name == "solfrac":
             # Bring these solfrac parameters to the subroutine as nvidia
             # does not permit offloaded kernels to access module parameters
             symtab = subroutine.symbol_table
-            symtab.add(symtab.lookup("pp_wgt"))
-            symtab.add(symtab.lookup("pp_len"))
+            if "pp_wgt" not in symtab:
+                symtab.add(symtab.lookup("pp_wgt"))
+            if "pp_len" not in symtab:
+                symtab.add(symtab.lookup("pp_len"))
 
         normalise_loops(
                 subroutine,

examples/nemo/scripts/utils.py

@@ -229,9 +229,6 @@ def normalise_loops(
         # Convert all array implicit loops to explicit loops
         explicit_loops = ArrayAssignment2LoopsTrans()
         for assignment in schedule.walk(Assignment):
-            # TODO #2951: Fix array assignments with dependencies
-            if schedule.name in ("fld_def",):
-                continue
             try:
                 explicit_loops.apply(
                     assignment, options={'verbose': True})

src/psyclone/psyir/nodes/array_mixin.py

@@ -645,7 +645,7 @@ def _get_effective_shape(self):
                 dtype = idx_expr.datatype
                 if isinstance(dtype, ArrayType):
                     # An array slice can be defined by a 1D slice of another
-                    # array, e.g. `a(b(1:4))` or `a(b)`.
+                    # array, e.g. `a(b(1:4))`, `a(b)` or a(b(:)).
                     indirect_array_shape = dtype.shape
                     if len(indirect_array_shape) > 1:
                         raise NotImplementedError(
@@ -655,7 +655,7 @@ def _get_effective_shape(self):
                             f"{len(indirect_array_shape)} dimensions.")
                     # pylint: disable=protected-access
                     if isinstance(idx_expr, ArrayMixin):
-                        shape.append(idx_expr._extent(idx))
+                        shape.append(idx_expr._get_effective_shape()[0])
                     else:
                         # We have some expression with a shape but no explicit
                         # indexing. The extent of this is then the SIZE of