introduce 'normalize_array_access' for prior execution to 'flatten_ar…

…rays' to allow for arrays start counting not with '1'

loki/transform/fortran_c_transform.py

@@ -12,7 +12,7 @@
 from loki.transform.transform_array_indexing import (
     shift_to_zero_indexing, invert_array_indices,
     resolve_vector_notation, normalize_range_indexing,
-    flatten_arrays
+    normalize_array_access, flatten_arrays
 )
 from loki.transform.transform_associates import resolve_associates
 from loki.transform.transform_utilities import (
@@ -37,7 +37,6 @@
 from loki.tools import as_tuple, flatten
 from loki.types import BasicType, DerivedType, SymbolAttributes
 
-
 __all__ = ['FortranCTransformation']
 
 
@@ -51,8 +50,9 @@ class FortranCTransformation(Transformation):
     # Set of standard module names that have no C equivalent
     __fortran_intrinsic_modules = ['ISO_FORTRAN_ENV', 'ISO_C_BINDING']
 
-    def __init__(self, header_modules=None, inline_elementals=True):
+    def __init__(self, header_modules=None, inline_elementals=True, order='F'):
         self.inline_elementals = inline_elementals
+        self.order = order
 
         # Maps from original type name to ISO-C and C-struct types
         self.c_structs = OrderedDict()
@@ -369,13 +369,14 @@ def generate_c_kernel(self, routine, **kwargs):
 
         # Clean up Fortran vector notation
         resolve_vector_notation(kernel)
+        normalize_array_access(kernel)
         normalize_range_indexing(kernel)
 
         # Convert array indexing to C conventions
         # TODO: Resolve reductions (eg. SUM(myvar(:)))
         invert_array_indices(kernel)
         shift_to_zero_indexing(kernel)
-        flatten_arrays(kernel, order="C", start_index=0)
+        flatten_arrays(kernel, order=self.order, start_index=0)
 
         # Inline all known parameters, since they can be used in declarations,
         # and thus need to be known before we can fetch them via getters.

loki/transform/transform_array_indexing.py

@@ -29,7 +29,7 @@
     'resolve_vector_notation', 'normalize_range_indexing',
     'promote_variables', 'promote_nonmatching_variables',
     'promotion_dimensions_from_loop_nest', 'demote_variables',
-    'flatten_arrays'
+    'flatten_arrays', 'normalize_array_access'
 ]
 
 
@@ -498,6 +498,43 @@ def demote_variables(routine, variable_names, dimensions):
 
     info(f'[Loki::Transform] Demoted variables in {routine.name}: {", ".join(variable_names)}')
 
+def normalize_array_access(routine):
+    """
+    Shift all arrays to start counting at "1"
+    """
+    def is_range_index(dim):
+        return isinstance(dim, sym.RangeIndex) and not dim.lower == 1
+
+    vmap = {}
+    for v in FindVariables(unique=False).visit(routine.body):
+        if isinstance(v, sym.Array):
+            new_dims = []
+            for i, d in enumerate(v.shape):
+                if isinstance(d, sym.RangeIndex):
+                    if isinstance(v.dimensions[i], sym.RangeIndex):
+                        start = simplify(v.dimensions[i].start - d.start + 1) if d.start is not None else None
+                        stop = simplify(v.dimensions[i].stop - d.start + 1) if d.stop is not None else None
+                        new_dims += [sym.RangeIndex((start, stop, d.step))]
+                    else:
+                        start = simplify(v.dimensions[i] - d.start + 1) if d.start is not None else None
+                        new_dims += [start]
+                else:
+                    new_dims += [v.dimensions[i]]
+            vmap[v] = v.clone(dimensions=as_tuple(new_dims))
+    routine.body = SubstituteExpressions(vmap).visit(routine.body)
+
+    vmap = {}
+    for v in routine.variables:
+        if isinstance(v, sym.Array):
+            new_dims = [sym.RangeIndex((1, simplify(d.upper - d.lower + 1)))
+                if is_range_index(d) else d for d in v.dimensions]
+            new_shape = [sym.RangeIndex((1, simplify(d.upper - d.lower + 1)))
+                if is_range_index(d) else d for d in v.shape]
+            new_type = v.type.clone(shape=as_tuple(new_shape))
+            vmap[v] = v.clone(dimensions=as_tuple(new_dims), type=new_type)
+    routine.variables = [vmap.get(v, v) for v in routine.variables]
+    normalize_range_indexing(routine)
+
 
 def flatten_arrays(routine, order='F', start_index=1):
     """
@@ -515,26 +552,29 @@ def flatten_arrays(routine, order='F', start_index=1):
     """
     def new_dims(dim, shape):
         if len(dim) > 1:
-            assert not isinstance(shape[-1], sym.RangeIndex)
-            _dim = [sym.Sum((dim[-2], sym.Product((shape[-2], dim[-1] - start_index))))]
+            if isinstance(shape[-2], sym.RangeIndex):
+                raise TypeError(f'Resolve shapes being of type RangeIndex, e.g., "{shape[-2]}" before flattening!')
+            _dim = (sym.Sum((dim[-2], sym.Product((shape[-2], dim[-1] - start_index)))),)
             new_dim = dim[:-2]
-            new_dim.extend(_dim)
+            new_dim += _dim
             return new_dims(new_dim, shape[:-1])
-        return as_tuple(dim)
+        return dim
 
-    assert order in ['F', 'C']
-    if order == 'C':
+    if order == 'F':
         array_map = {
-            var: var.clone(dimensions=new_dims(list(var.dimensions)[::-1], list(var.shape)[::-1]))
+            var: var.clone(dimensions=new_dims(var.dimensions[::-1], var.shape[::-1]))
             for var in FindVariables().visit(routine.body)
-                if isinstance(var, sym.Array) and var.shape and len(var.shape)
+            if isinstance(var, sym.Array) and var.shape and len(var.shape)
         }
-    elif order == 'F':
+    elif order == 'C':
         array_map = {
-            var: var.clone(dimensions=new_dims(list(var.dimensions), list(var.shape)))
+            var: var.clone(dimensions=new_dims(var.dimensions, var.shape))
             for var in FindVariables().visit(routine.body)
-                if isinstance(var, sym.Array) and var.shape and len(var.shape)
+            if isinstance(var, sym.Array) and var.shape and len(var.shape)
         }
+    else:
+        raise ValueError(f'Unsupported array order "{order}"')
+
     routine.body = SubstituteExpressions(array_map).visit(routine.body)
 
     routine.variables = [v.clone(dimensions=as_tuple(sym.Product(v.shape)),