diff --git a/loki/backend/cgen.py b/loki/backend/cgen.py index c7a476325..9e4880a4c 100644 --- a/loki/backend/cgen.py +++ b/loki/backend/cgen.py @@ -158,9 +158,8 @@ def visit_Subroutine(self, o, **kwargs): # Generate header with argument signature aptr = [] for a in o.arguments: - # TODO: Oh dear, the pointer derivation is beyond hacky; clean up! if isinstance(a, Array) > 0: - aptr += ['* restrict v_'] + aptr += ['* restrict '] elif isinstance(a.type.dtype, DerivedType): aptr += ['*'] elif a.type.pointer: @@ -176,18 +175,6 @@ def visit_Subroutine(self, o, **kwargs): # ...and generate the spec without imports and argument declarations body = [self.visit(o.spec, skip_imports=True, skip_argument_declarations=True, **kwargs)] - # Generate the array casts for pointer arguments - if any(isinstance(a, Array) for a in o.arguments): - body += [self.format_line('/* Array casts for pointer arguments */')] - for a in o.arguments: - if isinstance(a, Array): - dtype = self.visit(a.type, **kwargs) - # str(d).lower() is a bad hack to ensure caps-alignment - outer_dims = ''.join(f'[{self.visit(d, **kwargs).lower()}]' - for d in a.dimensions[1:]) - body += [self.format_line(dtype, ' (*', a.name.lower(), ')', outer_dims, ' = (', - dtype, ' (*)', outer_dims, ') v_', a.name.lower(), ';')] - # Fill the body body += [self.visit(o.body, **kwargs)] body += [self.format_line('return 0;')] diff --git a/loki/transform/fortran_c_transform.py b/loki/transform/fortran_c_transform.py index 37eb8fe56..45445a65e 100644 --- a/loki/transform/fortran_c_transform.py +++ b/loki/transform/fortran_c_transform.py @@ -11,7 +11,8 @@ from loki.transform.transformation import Transformation from loki.transform.transform_array_indexing import ( shift_to_zero_indexing, invert_array_indices, - resolve_vector_notation, normalize_range_indexing + resolve_vector_notation, normalize_array_shape_and_access, + flatten_arrays ) from loki.transform.transform_associates import resolve_associates from loki.transform.transform_utilities import ( @@ -36,7 +37,6 @@ from loki.tools import as_tuple, flatten from loki.types import BasicType, DerivedType, SymbolAttributes - __all__ = ['FortranCTransformation'] @@ -368,12 +368,13 @@ def generate_c_kernel(self, routine, **kwargs): # Clean up Fortran vector notation resolve_vector_notation(kernel) - normalize_range_indexing(kernel) + normalize_array_shape_and_access(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) # Inline all known parameters, since they can be used in declarations, # and thus need to be known before we can fetch them via getters. diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py index 946709ff1..bce200115 100644 --- a/loki/transform/transform_array_indexing.py +++ b/loki/transform/transform_array_indexing.py @@ -28,7 +28,8 @@ 'shift_to_zero_indexing', 'invert_array_indices', 'resolve_vector_notation', 'normalize_range_indexing', 'promote_variables', 'promote_nonmatching_variables', - 'promotion_dimensions_from_loop_nest', 'demote_variables' + 'promotion_dimensions_from_loop_nest', 'demote_variables', + 'flatten_arrays', 'normalize_array_shape_and_access' ] @@ -496,3 +497,86 @@ def demote_variables(routine, variable_names, dimensions): routine.spec = Transformer(decl_map).visit(routine.spec) info(f'[Loki::Transform] Demoted variables in {routine.name}: {", ".join(variable_names)}') + +def normalize_array_shape_and_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): + """ + Flatten arrays, converting multi-dimensional arrays to + one-dimensional arrays. + + Parameters + ---------- + routine : :any:`Subroutine` + The subroutine in which the variables should be promoted. + order : str + Assume Fortran (F) vs. C memory/array order. + start_index : int + Assume array indexing starts with `start_index`. + """ + def new_dims(dim, shape): + if len(dim) > 1: + 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 += _dim + return new_dims(new_dim, shape[:-1]) + return dim + + if order == 'C': + array_map = { + 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) + } + elif order == 'F': + array_map = { + 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) + } + 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)), + type=v.type.clone(shape=as_tuple(sym.Product(v.shape)))) + if isinstance(v, sym.Array) else v for v in routine.variables] diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py index ab48695e5..2ad912972 100644 --- a/tests/test_transform_array_indexing.py +++ b/tests/test_transform_array_indexing.py @@ -9,16 +9,26 @@ import pytest import numpy as np -from conftest import jit_compile, clean_test, available_frontends -from loki import Subroutine +from conftest import jit_compile, jit_compile_lib, clean_test, available_frontends +from loki import Subroutine, FindVariables, Array from loki.expression import symbols as sym -from loki.transform import promote_variables, demote_variables, normalize_range_indexing - +from loki.transform import ( + promote_variables, demote_variables, normalize_range_indexing, + invert_array_indices, flatten_arrays, + normalize_array_shape_and_access + ) +from loki.transform import ( + FortranCTransformation + ) +from loki.build import Builder @pytest.fixture(scope='module', name='here') def fixture_here(): return Path(__file__).parent +@pytest.fixture(scope='module', name='builder') +def fixture_builder(here): + return Builder(source_dirs=here, build_dir=here/'build') @pytest.mark.parametrize('frontend', available_frontends()) def test_transform_promote_variable_scalar(here, frontend): @@ -315,3 +325,215 @@ def test_transform_demote_dimension_arguments(here, frontend): assert np.all(vec1 == 3) and np.sum(vec1) == 3 assert np.all(vec2 == 2) and np.sum(vec2) == 6 assert np.all(matrix == 16) and np.sum(matrix) == 32 + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('start_index', (0, 1, 5)) +def test_transform_normalize_array_shape_and_access(here, frontend, start_index): + """ + Test flattening or arrays, meaning converting multi-dimensional + arrays to one-dimensional arrays including corresponding + index arithmetic. + """ + fcode = f""" + subroutine transform_normalize_array_shape_and_access(x1, x2, x3, x4, l1, l2, l3, l4) + implicit none + integer :: i1, i2, i3, i4, c1, c2, c3, c4 + integer, intent(in) :: l1, l2, l3, l4 + integer, intent(inout) :: x1({start_index}:l1+{start_index}-1) + integer, intent(inout) :: x2({start_index}:l2+{start_index}-1, & + & {start_index}:l1+{start_index}-1) + integer, intent(inout) :: x3({start_index}:l3+{start_index}-1, & + & {start_index}:l2+{start_index}-1, {start_index}:l1+{start_index}-1) + integer, intent(inout) :: x4({start_index}:l4+{start_index}-1, & + & {start_index}:l3+{start_index}-1, {start_index}:l2+{start_index}-1, & + & {start_index}:l1+{start_index}-1) + c1 = 1 + c2 = 1 + c3 = 1 + c4 = 1 + x1({start_index}:l4+{start_index}-1) = 0 + do i1={start_index},l1+{start_index}-1 + x1(i1) = c1 + do i2={start_index},l2+{start_index}-1 + x2(i2, i1) = c2*10 + c1 + do i3={start_index},l3+{start_index}-1 + x3(i3, i2, i1) = c3*100 + c2*10 + c1 + do i4={start_index},l4+{start_index}-1 + x4(i4, i3, i2, i1) = c4*1000 + c3*100 + c2*10 + c1 + c4 = c4 + 1 + end do + c3 = c3 + 1 + end do + c2 = c2 + 1 + end do + c1 = c1 + 1 + end do + + end subroutine transform_normalize_array_shape_and_access + """ + def init_arguments(l1, l2, l3, l4): + x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32) + x2 = np.zeros(shape=(l2,l1,), order='F', dtype=np.int32) + x3 = np.zeros(shape=(l3,l2,l1,), order='F', dtype=np.int32) + x4 = np.zeros(shape=(l4,l3,l2,l1,), order='F', dtype=np.int32) + return x1, x2, x3, x4 + + def validate_routine(routine): + arrays = [var for var in FindVariables().visit(routine.body) if isinstance(var, Array)] + for arr in arrays: + assert all(not isinstance(shape, sym.RangeIndex) for shape in arr.shape) + + l1 = 2 + l2 = 3 + l3 = 4 + l4 = 5 + routine = Subroutine.from_source(fcode, frontend=frontend) + normalize_range_indexing(routine) # Fix OMNI nonsense + filepath = here/(f'{routine.name}_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + orig_x1, orig_x2, orig_x3, orig_x4 = init_arguments(l1, l2, l3, l4) + function(orig_x1, orig_x2, orig_x3, orig_x4, l1, l2, l3, l4) + clean_test(filepath) + + routine = Subroutine.from_source(fcode, frontend=frontend) + normalize_array_shape_and_access(routine) + normalize_range_indexing(routine) + filepath = here/(f'{routine.name}_normalized_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + x1, x2, x3, x4 = init_arguments(l1, l2, l3, l4) + function(x1, x2, x3, x4, l1, l2, l3, l4) + validate_routine(routine) + clean_test(filepath) + + assert (x1 == orig_x1).all() + assert (x2 == orig_x2).all() + assert (x3 == orig_x3).all() + assert (x4 == orig_x4).all() + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('start_index', (0, 1, 5)) +def test_transform_flatten_arrays(here, frontend, builder, start_index): + """ + Test flattening or arrays, meaning converting multi-dimensional + arrays to one-dimensional arrays including corresponding + index arithmetic. + """ + fcode = f""" + subroutine transform_flatten_arrays(x1, x2, x3, x4, l1, l2, l3, l4) + implicit none + integer :: i1, i2, i3, i4, c1, c2, c3, c4 + integer, intent(in) :: l1, l2, l3, l4 + integer, intent(inout) :: x1({start_index}:l1+{start_index}-1) + integer, intent(inout) :: x2({start_index}:l2+{start_index}-1, & + & {start_index}:l1+{start_index}-1) + integer, intent(inout) :: x3({start_index}:l3+{start_index}-1, & + & {start_index}:l2+{start_index}-1, {start_index}:l1+{start_index}-1) + integer, intent(inout) :: x4({start_index}:l4+{start_index}-1, & + & {start_index}:l3+{start_index}-1, {start_index}:l2+{start_index}-1, & + & {start_index}:l1+{start_index}-1) + c1 = 1 + c2 = 1 + c3 = 1 + c4 = 1 + do i1={start_index},l1+{start_index}-1 + x1(i1) = c1 + do i2={start_index},l2+{start_index}-1 + x2(i2, i1) = c2*10 + c1 + do i3={start_index},l3+{start_index}-1 + x3(i3, i2, i1) = c3*100 + c2*10 + c1 + do i4={start_index},l4+{start_index}-1 + x4(i4, i3, i2, i1) = c4*1000 + c3*100 + c2*10 + c1 + c4 = c4 + 1 + end do + c3 = c3 + 1 + end do + c2 = c2 + 1 + end do + c1 = c1 + 1 + end do + + end subroutine transform_flatten_arrays + """ + def init_arguments(l1, l2, l3, l4, flattened=False): + x1 = np.zeros(shape=(l1,), order='F', dtype=np.int32) + x2 = np.zeros(shape=(l2*l1) if flattened else (l2,l1,), order='F', dtype=np.int32) + x3 = np.zeros(shape=(l3*l2*l1) if flattened else (l3,l2,l1,), order='F', dtype=np.int32) + x4 = np.zeros(shape=(l4*l3*l2*l1) if flattened else (l4,l3,l2,l1,), order='F', dtype=np.int32) + return x1, x2, x3, x4 + + def validate_routine(routine): + arrays = [var for var in FindVariables().visit(routine.body) if isinstance(var, Array)] + assert all(len(arr.dimensions) == 1 for arr in arrays) + assert all(len(arr.shape) == 1 for arr in arrays) + + l1 = 2 + l2 = 3 + l3 = 4 + l4 = 5 + # Test the original implementation + routine = Subroutine.from_source(fcode, frontend=frontend) + normalize_range_indexing(routine) # Fix OMNI nonsense + filepath = here/(f'{routine.name}_{start_index}_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + orig_x1, orig_x2, orig_x3, orig_x4 = init_arguments(l1, l2, l3, l4) + function(orig_x1, orig_x2, orig_x3, orig_x4, l1, l2, l3, l4) + clean_test(filepath) + + # Test flattening order='F' + f_routine = Subroutine.from_source(fcode, frontend=frontend) + normalize_array_shape_and_access(f_routine) + normalize_range_indexing(f_routine) # Fix OMNI nonsense + flatten_arrays(routine=f_routine, order='F', start_index=1) + filepath = here/(f'{f_routine.name}_{start_index}_flattened_F_{frontend}.f90') + function = jit_compile(f_routine, filepath=filepath, objname=routine.name) + f_x1, f_x2, f_x3, f_x4 = init_arguments(l1, l2, l3, l4, flattened=True) + function(f_x1, f_x2, f_x3, f_x4, l1, l2, l3, l4) + validate_routine(f_routine) + clean_test(filepath) + + assert (f_x1 == orig_x1.flatten(order='F')).all() + assert (f_x2 == orig_x2.flatten(order='F')).all() + assert (f_x3 == orig_x3.flatten(order='F')).all() + assert (f_x4 == orig_x4.flatten(order='F')).all() + + # Test flattening order='C' + c_routine = Subroutine.from_source(fcode, frontend=frontend) + normalize_array_shape_and_access(c_routine) + normalize_range_indexing(c_routine) # Fix OMNI nonsense + invert_array_indices(c_routine) + flatten_arrays(routine=c_routine, order='C', start_index=1) + filepath = here/(f'{c_routine.name}_{start_index}_flattened_C_{frontend}.f90') + function = jit_compile(c_routine, filepath=filepath, objname=routine.name) + c_x1, c_x2, c_x3, c_x4 = init_arguments(l1, l2, l3, l4, flattened=True) + function(c_x1, c_x2, c_x3, c_x4, l1, l2, l3, l4) + validate_routine(c_routine) + clean_test(filepath) + + assert f_routine.body == c_routine.body + + assert (c_x1 == orig_x1.flatten(order='F')).all() + assert (c_x2 == orig_x2.flatten(order='F')).all() + assert (c_x3 == orig_x3.flatten(order='F')).all() + assert (c_x4 == orig_x4.flatten(order='F')).all() + + # Test C transpilation (which includes flattening) + f2c_routine = Subroutine.from_source(fcode, frontend=frontend) + f2c = FortranCTransformation() + f2c.apply(source=f2c_routine, path=here) + libname = f'fc_{f2c_routine.name}_{start_index}_{frontend}' + c_kernel = jit_compile_lib([f2c.wrapperpath, f2c.c_path], path=here, name=libname, builder=builder) + fc_function = c_kernel.transform_flatten_arrays_fc_mod.transform_flatten_arrays_fc + f2c_x1, f2c_x2, f2c_x3, f2c_x4 = init_arguments(l1, l2, l3, l4, flattened=True) + fc_function(f2c_x1, f2c_x2, f2c_x3, f2c_x4, l1, l2, l3, l4) + validate_routine(c_routine) + + assert (f2c_x1 == orig_x1.flatten(order='F')).all() + assert (f2c_x2 == orig_x2.flatten(order='F')).all() + assert (f2c_x3 == orig_x3.flatten(order='F')).all() + assert (f2c_x4 == orig_x4.flatten(order='F')).all() + + builder.clean() + clean_test(filepath) + f2c.wrapperpath.unlink() + f2c.c_path.unlink() diff --git a/tests/test_transpile.py b/tests/test_transpile.py index c4a9a2b43..17235bad9 100644 --- a/tests/test_transpile.py +++ b/tests/test_transpile.py @@ -940,7 +940,8 @@ def test_transpile_expressions(here, builder, frontend): # Make sure minus signs are represented correctly in the C code ccode = cgen(routine) - assert 'vector[i - 1 - 1]' in ccode # double minus due to index shift to 0 + # double minus due to index shift to 0 + assert 'vector[i - 1 - 1]' in ccode or 'vector[-1 + i - 1]' in ccode assert 'vector[i - 1]' in ccode assert '-scalar' in ccode # scalar with negative sign