From 924ae71763ad949888af70bb98f39016565c3d86 Mon Sep 17 00:00:00 2001 From: Michael Staneker Date: Sat, 5 Oct 2024 12:11:37 +0200 Subject: [PATCH] resolve_vector_notation: simplification and avoid unnecessary entry in symbol table --- loki/transformations/array_indexing.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/loki/transformations/array_indexing.py b/loki/transformations/array_indexing.py index 9d607675a..cdbfb4267 100644 --- a/loki/transformations/array_indexing.py +++ b/loki/transformations/array_indexing.py @@ -159,7 +159,7 @@ def resolve_vector_notation(routine): vmap = {} # find available loops and create map {(lower, upper, step): loop_variable} loops = FindNodes(Loop).visit(routine.body) - loop_map = {(loop.bounds.lower, loop.bounds.upper, loop.bounds.step if loop.bounds.step is not None else 1): + loop_map = {(loop.bounds.lower, loop.bounds.upper, loop.bounds.step or 1): loop.variable for loop in loops} for stmt in FindNodes(Assignment).visit(routine.body): # Loop over all variables and replace them with loop indices @@ -177,15 +177,17 @@ def resolve_vector_notation(routine): ivar_basename = f'i_{stmt.lhs.basename}' for i, dim, s in zip(count(), v.dimensions, as_tuple(v.shape)): if isinstance(dim, sym.RangeIndex): - # Create new index variable - vtype = SymbolAttributes(BasicType.INTEGER) - ivar = sym.Variable(name=f'{ivar_basename}_{i}', type=vtype, scope=routine) # create tuple to test whether an appropriate loop is already available test_range = (sym.IntLiteral(1), s, 1) if not isinstance(s, sym.RangeIndex)\ else (s.lower, s.upper, 1) # actually test for it if test_range in loop_map: + # Use index variable of available matching loop ivar = loop_map[test_range] + else: + # Create new index variable + vtype = SymbolAttributes(BasicType.INTEGER) + ivar = sym.Variable(name=f'{ivar_basename}_{i}', type=vtype, scope=routine) shape_index_map[(i, s)] = ivar index_range_map[ivar] = s