From f8e813cb5f53cd0a94a12004d6b565a011be853e Mon Sep 17 00:00:00 2001 From: Michael Staneker Date: Wed, 21 Feb 2024 14:48:10 +0000 Subject: [PATCH 1/3] [F2C] 'shift_to_zero_indexing' allow to ignore dimensions --- loki/transform/transform_array_indexing.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py index cadecb792..77afd1578 100644 --- a/loki/transform/transform_array_indexing.py +++ b/loki/transform/transform_array_indexing.py @@ -33,7 +33,7 @@ ] -def shift_to_zero_indexing(routine): +def shift_to_zero_indexing(routine, ignore=()): """ Shift all array indices to adjust to 0-based indexing conventions (eg. for C or Python) """ @@ -47,7 +47,14 @@ def shift_to_zero_indexing(routine): # no shift for stop because Python ranges are [start, stop) new_dims += [sym.RangeIndex((start, d.stop, d.step))] else: - new_dims += [d - sym.Literal(1)] + shift = True + for var in FindVariables().visit(d): + if var in ignore: + shift = False + if shift: + new_dims += [d - sym.Literal(1)] + else: + new_dims += [d] vmap[v] = v.clone(dimensions=as_tuple(new_dims)) routine.body = SubstituteExpressions(vmap).visit(routine.body) From e94a942e288ae32fc599cfee02f380e184abeb92 Mon Sep 17 00:00:00 2001 From: Michael Staneker Date: Mon, 4 Mar 2024 12:11:10 +0000 Subject: [PATCH 2/3] [F2C] continued: 'shift_to_zero_indexing' allow to ignore dimensions including unit testing --- loki/transform/transform_array_indexing.py | 7 +++ tests/test_transform_array_indexing.py | 60 +++++++++++++++++++++- 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py index 77afd1578..50238ea91 100644 --- a/loki/transform/transform_array_indexing.py +++ b/loki/transform/transform_array_indexing.py @@ -36,6 +36,13 @@ def shift_to_zero_indexing(routine, ignore=()): """ Shift all array indices to adjust to 0-based indexing conventions (eg. for C or Python) + + Parameters + ---------- + routine : :any:`Subroutine` + The subroutine in which the array dimensions should be shifted + ignore : list of str + Dimensions (or rather variables being dimensions) to be ignored """ vmap = {} for v in FindVariables(unique=False).visit(routine.body): diff --git a/tests/test_transform_array_indexing.py b/tests/test_transform_array_indexing.py index 2cd9cbaef..dcfaf58e3 100644 --- a/tests/test_transform_array_indexing.py +++ b/tests/test_transform_array_indexing.py @@ -10,12 +10,12 @@ import numpy as np from conftest import jit_compile, jit_compile_lib, clean_test, available_frontends -from loki import Module, Subroutine, FindVariables, Array +from loki import Module, Subroutine, FindVariables, Array, fgen from loki.expression import symbols as sym from loki.transform import ( promote_variables, demote_variables, normalize_range_indexing, invert_array_indices, flatten_arrays, - normalize_array_shape_and_access + normalize_array_shape_and_access, shift_to_zero_indexing ) from loki.transform import ( FortranCTransformation @@ -572,3 +572,59 @@ def validate_routine(routine): clean_test(filepath) f2c.wrapperpath.unlink() f2c.c_path.unlink() + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('ignore', ((), ('i2',), ('i4', 'i1'))) +def test_shift_to_zero_indexing(frontend, ignore): + """ + Test shifting array dimensions to zero (or rather shift dimension `dim` + to `dim - 1`). This does not produce valid Fortran, but is part of the + F2C transpilation logic. + """ + fcode = """ + subroutine transform_shift_indexing(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(l1) + integer, intent(inout) :: x2(l2, l1) + integer, intent(inout) :: x3(l3, l2, l1) + integer, intent(inout) :: x4(l4, l3, l2, l1) + c1 = 1 + c2 = 1 + c3 = 1 + c4 = 1 + do i1=1,l1 + x1(i1) = c1 + do i2=1,l2 + x2(i2, i1) = c2*10 + c1 + do i3=1,l3 + x3(i3, i2, i1) = c3*100 + c2*10 + c1 + do i4=1,l4 + 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_shift_indexing + """ + + expected_dims = {'x1': ('i1',), 'x2': ('i2', 'i1'), + 'x3': ('i3', 'i2', 'i1'), 'x4': ('i4', 'i3', 'i2', 'i1')} + routine = Subroutine.from_source(fcode, frontend=frontend) + arrays = [var for var in FindVariables().visit(routine.body) if isinstance(var, sym.Array)] + for array in arrays: + assert array.dimensions == expected_dims[array.name] + + shift_to_zero_indexing(routine, ignore=ignore) + + arrays = [var for var in FindVariables().visit(routine.body) if isinstance(var, sym.Array)] + for array in arrays: + dimensions = tuple(sym.Sum((sym.Scalar(name=dim), sym.Product((-1, sym.IntLiteral(1))))) + if dim not in ignore else dim for dim in expected_dims[array.name]) + assert fgen(array.dimensions) == fgen(dimensions) From 12aaeb0723ebb0bdf3d0fe3e85baf137d04038b2 Mon Sep 17 00:00:00 2001 From: Michael Staneker Date: Tue, 5 Mar 2024 14:02:50 +0000 Subject: [PATCH 3/3] 'shift_to_zero_indexing' some streamlining and improvements --- loki/transform/transform_array_indexing.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/loki/transform/transform_array_indexing.py b/loki/transform/transform_array_indexing.py index 50238ea91..060451b93 100644 --- a/loki/transform/transform_array_indexing.py +++ b/loki/transform/transform_array_indexing.py @@ -33,7 +33,7 @@ ] -def shift_to_zero_indexing(routine, ignore=()): +def shift_to_zero_indexing(routine, ignore=None): """ Shift all array indices to adjust to 0-based indexing conventions (eg. for C or Python) @@ -42,8 +42,10 @@ def shift_to_zero_indexing(routine, ignore=()): routine : :any:`Subroutine` The subroutine in which the array dimensions should be shifted ignore : list of str - Dimensions (or rather variables being dimensions) to be ignored + List of variable names for which, if found in the dimension expression + of an array subscript, that dimension is not shifted to zero. """ + ignore = as_tuple(ignore) vmap = {} for v in FindVariables(unique=False).visit(routine.body): if isinstance(v, sym.Array): @@ -54,14 +56,10 @@ def shift_to_zero_indexing(routine, ignore=()): # no shift for stop because Python ranges are [start, stop) new_dims += [sym.RangeIndex((start, d.stop, d.step))] else: - shift = True - for var in FindVariables().visit(d): - if var in ignore: - shift = False - if shift: - new_dims += [d - sym.Literal(1)] - else: + if ignore and any(var in ignore for var in FindVariables().visit(d)): new_dims += [d] + else: + new_dims += [d - sym.Literal(1)] vmap[v] = v.clone(dimensions=as_tuple(new_dims)) routine.body = SubstituteExpressions(vmap).visit(routine.body)