Merge pull request #414 from ecmwf-ifs/naml-dimension-range-step

Dimension: Support stepping, implicit aliases and remove contrainsts

loki/dimension.py

@@ -18,43 +18,74 @@ class Dimension:
     ----------
     name : string
         Name of the dimension to identify in configurations
-    index : string
+    index : string or tuple of str
         String representation of the predominant loop index variable
-        associated with this dimension.
-    size : string
+        associated with this dimension; can be one or several.
+    size : string or tuple of str
         String representation of the predominant size variable used
-        to declare array shapes using this dimension.
+        to declare array shapes; can be one or several.
+    lower : string or tuple of str
+        String representation of the lower bound variable used to
+        declare iteration spaces; can be one or several.
+    lower : string or tuple of str
+        String representation of the upper bound variable used to
+        declare iteration spaces; can be one or several.
     bounds : tuple of strings
         String representation of the variables usually used to denote
         the iteration bounds of this dimension.
+
+        **WARNING:** This argument is deprecated, instead ``lower``
+        and ``upper`` should be used.
     aliases : list or tuple of strings
         String representations of alternative size variables that are
         used to define arrays shapes of this dimension (eg. alternative
         names used in "driver" subroutines).
+
+        **WARNING:** This argument is deprecated, instead a tuple of
+        variables names should be provided for ``size``.
     bounds_aliases : list or tuple of strings
         String representations of alternative bounds variables that are
         used to define loop ranges.
+
+        **WARNING:** This argument is deprecated, instead a tuple of
+        variables names should be provided for ``lower`` and
+        ``upper``.
     index_aliases : list or tuple of strings
         String representations of alternative loop index variables associated
         with this dimension.
+
+        **WARNING:** This argument is deprecated, instead a tuple of
+        variables names should be provided for ``index``.
     """
 
-    def __init__(self, name=None, index=None, bounds=None, size=None, aliases=None,
-                 bounds_aliases=None, index_aliases=None):
+    def __init__(
+            self, name=None, index=None, size=None, lower=None,
+            upper=None, step=None, aliases=None, bounds=None,
+            bounds_aliases=None, index_aliases=None
+    ):
         self.name = name
-        self._index = index
-        self._bounds = as_tuple(bounds)
-        self._size = size
-        self._aliases = as_tuple(aliases)
-        self._index_aliases = as_tuple(index_aliases)
 
+        if bounds:
+            # Backward compat for ``bounds`` contructor argument
+            assert not lower and not upper and len(bounds) == 2
+            lower = (bounds[0],)
+            upper = (bounds[1],)
+
+        # Store one or more strings for dimension variables
+        self._index = as_tuple(index) or None
+        self._size = as_tuple(size) or None
+        self._lower = as_tuple(lower) or None
+        self._upper = as_tuple(upper) or None
+        self._step = as_tuple(step) or None
+
+        # Keep backward-compatibility for constructor arguments
+        if aliases:
+            self._size += as_tuple(aliases)
+        if index_aliases:
+            self._index += as_tuple(index_aliases)
         if bounds_aliases:
-            if len(bounds_aliases) != 2:
-                raise RuntimeError(f'Start and end both needed for horizontal bounds aliases in {self.name}')
-            if bounds_aliases[0].split('%')[0] != bounds_aliases[1].split('%')[0]:
-                raise RuntimeError(f'Inconsistent root name for horizontal bounds aliases in {self.name}')
-
-        self._bounds_aliases = as_tuple(bounds_aliases)
+            self._lower = as_tuple(self._lower) + (bounds_aliases[0],)
+            self._upper = as_tuple(self._upper) + (bounds_aliases[1],)
 
     def __repr__(self):
         """ Pretty-print dimension details """
@@ -69,68 +100,92 @@ def variables(self):
         return (self.index, self.size) + self.bounds
 
     @property
-    def size(self):
+    def sizes(self):
         """
-        String that matches the size expression of a data space (variable allocation).
+        Tuple of strings that match the primary size and all secondary size expressions.
+
+        .. note::
+            For derived expressions, like ``end - start + 1`` or
+            ``1:size``, please use :any:`size_expressions`.
         """
         return self._size
 
     @property
-    def index(self):
+    def size(self):
         """
-        String that matches the primary index expression of an iteration space (loop).
+        String that matches the primary size expression of a data space (variable allocation).
         """
-        return self._index
+        return self.sizes[0] if self.sizes else None
 
     @property
-    def bounds(self):
+    def indices(self):
         """
-        Tuple of expression string that represent the bounds of an iteration space.
+        Tuple of strings that matche the primary index and all secondary index expressions.
         """
-        return self._bounds
+        return self._index
 
     @property
-    def range(self):
+    def index(self):
         """
-        String that matches the range expression of an iteration space (loop).
+        String that matches the primary index expression of an iteration space (loop).
         """
-        return f'{self._bounds[0]}:{self._bounds[1]}'
+        return self.indices[0] if self.indices else None
 
     @property
-    def size_expressions(self):
+    def lower(self):
         """
-        A list of all expression strings representing the size of a data space.
+        String or tuple of strings that matches the lower bound of the iteration space.
+        """
+        return self._lower[0] if self._lower and len(self._lower) == 1 else self._lower
 
-        This includes generic aliases, like ``end - start + 1`` or ``1:size`` ranges.
+    @property
+    def upper(self):
         """
-        exprs = as_tuple(self.size)
-        if self._aliases:
-            exprs += as_tuple(self._aliases)
-        exprs += (f'1:{self.size}', )
-        if self._bounds:
-            exprs += (f'{self._bounds[1]} - {self._bounds[0]} + 1', )
-        return exprs
+        String or tuple of strings that matches the upper bound of the iteration space.
+        """
+        return self._upper[0] if self._upper and len(self._upper) == 1 else self._upper
 
     @property
-    def bounds_expressions(self):
+    def step(self):
         """
-        A list of all expression strings representing the bounds of a data space.
+        String or tuple of strings that matches the step size of the iteration space.
         """
+        return self._step[0] if self._step and len(self._step) == 1 else self._step
 
-        exprs = [(b,) for b in self.bounds]
-        if self._bounds_aliases:
-            exprs = [expr + (b,) for expr, b in zip(exprs, self._bounds_aliases)]
+    @property
+    def bounds(self):
+        """
+        Tuple of expression string that represent the bounds of an iteration space.
 
-        return as_tuple(exprs)
+        .. note::
+            If mutiple lower or upper bound string have been provided,
+            only the first pair will be used.
+        """
+        return (
+            self.lower[0] if isinstance(self.lower, tuple) else self.lower,
+            self.upper[0] if isinstance(self.upper, tuple) else self.upper
+        )
 
     @property
-    def index_expressions(self):
+    def range(self):
         """
-        A list of all expression strings representing the index expression of an iteration space (loop).
+        String that matches the range expression of an iteration space (loop).
+
+        .. note::
+            If mutiple lower or upper bound string have been provided,
+            only the first pair will be used.
         """
+        return f'{self.bounds[0]}:{self.bounds[1]}'
 
-        exprs = [self.index,]
-        if self._index_aliases:
-            exprs += list(self._index_aliases)
+    @property
+    def size_expressions(self):
+        """
+        A list of all expression strings representing the size of a data space.
 
-        return as_tuple(exprs)
+        This includes generic aliases, like ``end - start + 1`` or ``1:size`` ranges.
+        """
+        exprs = self.sizes
+        exprs += (f'1:{self.size}', )
+        if self.bounds:
+            exprs += (f'{self.bounds[1]} - {self.bounds[0]} + 1', )
+        return exprs

loki/tests/test_dimension.py

@@ -8,7 +8,51 @@
 import pytest
 
 from loki import Subroutine, Dimension, FindNodes, Loop
+from loki.batch import SchedulerConfig
+from loki.expression import symbols as sym
 from loki.frontend import available_frontends
+from loki.scope import Scope, SymbolAttributes
+from loki.types import BasicType
+
+
+def test_dimension_properties():
+    """
+    Test that :any:`Dimension` objects store the correct strings.
+    """
+    scope = Scope()
+    type_int = SymbolAttributes(dtype=BasicType.INTEGER)
+    i = sym.Variable(name='i', type=type_int, scope=scope)
+    n = sym.Variable(name='n', type=type_int, scope=scope)
+    z = sym.Variable(name='z', type=type_int, scope=scope)
+    one = sym.IntLiteral(1)
+    two = sym.IntLiteral(2)
+
+    simple = Dimension('simple', index='i', upper='n', size='z')
+    assert simple.index == i
+    assert simple.upper == n
+    assert simple.size == z
+
+    detail = Dimension(index='i', lower='1', upper='n', step='2', size='z')
+    assert detail.index == i
+    assert detail.lower == one
+    assert detail.upper == n
+    assert detail.step == two
+    assert detail.size == z
+    # Check derived properties
+    assert detail.bounds == (one, n)
+    assert detail.range == sym.LoopRange((1, n))
+
+    multi = Dimension(
+        index=('i', 'idx'), lower=('1', 'start'), upper=('n', 'end'), size='z'
+    )
+    assert multi.index == i
+    assert multi.indices == (i, sym.Variable(name='idx', type=type_int, scope=scope))
+    assert multi.lower == (one, sym.Variable(name='start', type=type_int, scope=scope))
+    assert multi.upper == (n, sym.Variable(name='end', type=type_int, scope=scope))
+    assert multi.size == z
+    # Check derived properties
+    assert multi.bounds ==  (one, n)
+    assert multi.range == sym.LoopRange((1, n))
 
 
 @pytest.mark.parametrize('frontend', available_frontends())
@@ -67,7 +111,48 @@ def test_dimension_index_range(frontend):
 
     # Test the correct creation of horizontal dim with aliased bounds vars
     _ = Dimension('test_dim_alias', bounds_aliases=('bnds%start', 'bnds%end'))
-    with pytest.raises(RuntimeError):
-        _ = Dimension('test_dim_alias', bounds_aliases=('bnds%start',))
-    with pytest.raises(RuntimeError):
-        _ = Dimension('test_dim_alias', bounds_aliases=('bnds%start', 'some_other_bnds%end'))
+
+
+def test_dimension_config(tmp_path):
+    """
+    Test that :any:`Dimension` objects get created from
+    :any:`SchedulerConfig` correctly.
+    """
+    scope = Scope()
+    type_int = SymbolAttributes(dtype=BasicType.INTEGER)
+    type_deferred = SymbolAttributes(dtype=BasicType.DEFERRED)
+    ibl = sym.Variable(name='ibl', type=type_int, scope=scope)
+    nblocks = sym.Variable(name='nblocks', type=type_int, scope=scope)
+    start = sym.Variable(name='start', type=type_int, scope=scope)
+    end = sym.Variable(name='end', type=type_int, scope=scope)
+    dim = sym.Variable(name='dim', type=type_deferred, scope=scope)
+    one = sym.IntLiteral(1)
+
+    config_str = """
+[dimensions.dim_a]
+  size = 'NBLOCKS'
+  index = 'IBL'
+  bounds = ['START', 'END']
+  aliases = ['DIM%START', 'DIM%END']
+
+[dimensions.dim_b]
+  size = 'nblocks'
+  index = 'ibl'
+  lower = ['1', 'start', 'dim%start']
+  upper = ['nblocks', 'end', 'dim%end']
+"""
+    cfg_path = tmp_path/'test_config.yml'
+    cfg_path.write_text(config_str)
+
+    config = SchedulerConfig.from_file(cfg_path)
+    dim_a = config.dimensions['dim_a']
+    assert dim_a.size == nblocks
+    assert dim_a.index == ibl
+    assert dim_a.bounds == (start, end)
+
+    dim_b = config.dimensions['dim_b']
+    assert dim_b.size == nblocks
+    assert dim_b.index == ibl
+    assert dim_b.bounds == (sym.IntLiteral(1), nblocks)
+    assert dim_b.lower == (one, start, start.clone(parent=dim))
+    assert dim_b.upper == (nblocks, end, end.clone(parent=dim))

loki/transformations/block_index_transformations.py

@@ -216,7 +216,7 @@ def process_body(self, body, definitions, successors, targets, exclude_arrays):
             if call.name in targets:
                 _args = {a: d for d, a in call.arg_map.items() if isinstance(d, Array)}
                 _vars += [a for a, d in _args.items()
-                         if any(v in d.shape for v in self.horizontal.size_expressions) and a.parents]
+                          if any(v in d.shape for v in self.horizontal.sizes) and a.parents]
 
         # replace per-block view pointers with full field pointers
         vmap = {var: var.clone(name=var.name_parts[-1] + '_FIELD',
@@ -375,7 +375,7 @@ def get_block_index(self, routine):
         variable_map = routine.variable_map
         if (block_index := variable_map.get(self.block_dim.index, None)):
             return block_index
-        if (block_index := [i for i in self.block_dim.index_expressions
+        if (block_index := [i for i in self.block_dim.indices
                             if i.split('%', maxsplit=1)[0] in variable_map]):
             return routine.resolve_typebound_var(block_index[0], variable_map)
         return None
@@ -690,8 +690,7 @@ def process_driver(self, routine, targets):
         with pragma_regions_attached(routine):
             with pragmas_attached(routine, ir.Loop):
                 loops = FindNodes(ir.Loop).visit(routine.body)
-                loops = [loop for loop in loops if loop.variable == self.block_dim.index
-                        or loop.variable in self.block_dim._index_aliases]
+                loops = [loop for loop in loops if loop.variable in self.block_dim.indices]
 
                 # Remove parallel regions around block loops
                 pragma_region_map = {}

loki/transformations/pool_allocator.py

@@ -626,7 +626,7 @@ def _create_stack_allocation(self, stack_ptr, stack_end, ptr_var, arr, stack_siz
 
         # If the array size is not a multiple of NPROMA, then we pad the allocation to avoid
         # potential alignment issues on device
-        if not self.horizontal or not any(s in dim for s in self.horizontal.size_expressions):
+        if not self.horizontal or not any(s in dim for s in self.horizontal.sizes):
             arr_type_bytes = InlineCall(function=Variable(name='MAX'),
                         parameters=(arr_type_bytes, Literal(8)), kw_parameters=())
         if self.cray_ptr_loc_rhs:

loki/transformations/single_column/claw.py

@@ -155,7 +155,7 @@ def hoist_dimension_from_call(self, caller, wrap=True):
         wrap : bool
             Flag to trigger the creation of a loop to wrap the kernel call in.
         """
-        size_expressions = self.horizontal.size_expressions
+        size_expressions = self.horizontal.sizes
         replacements = {}
 
         for call in FindNodes(CallStatement).visit(caller.body):
@@ -274,7 +274,7 @@ def transform_subroutine(self, routine, **kwargs):
 
         # Store the names of all variables that we are about to remove
         claw_vars = [v.name for v in routine.variables
-                     if isinstance(v, Array) and v.shape[0] in self.horizontal.size_expressions]
+                     if isinstance(v, Array) and v.shape[0] in self.horizontal.sizes]
 
         # The CLAW assumes that variables defining dimension sizes or iteration spaces
         # exist in both driver and kernel as local variables. We often rely on implicit