#2661 rename index arg to set_upper_bound to halo_depth

src/psyclone/domain/lfric/lfric_loop.py

@@ -49,7 +49,7 @@
 from psyclone.f2pygen import CallGen, CommentGen
 from psyclone.psyGen import InvokeSchedule, HaloExchange
 from psyclone.psyir.backend.fortran import FortranWriter
-from psyclone.psyir.nodes import (ArrayReference, ACCRegionDirective,
+from psyclone.psyir.nodes import (ArrayReference, ACCRegionDirective, DataNode,
                                   Loop, Literal, OMPRegionDirective, Reference,
                                   Routine, Schedule)
 from psyclone.psyir.symbols import DataSymbol, INTEGER_TYPE
@@ -244,7 +244,7 @@ def load(self, kern):
                     # first halo cell starts immediately after the last owned
                     # cell, and the cell indices are contiguous.
                     self.set_lower_bound("cell_halo_start")
-                self.set_upper_bound("cell_halo", index=kern.halo_depth)
+                self.set_upper_bound("cell_halo", halo_depth=kern.halo_depth)
                 return
 
         if not Config.get().distributed_memory:
@@ -256,7 +256,7 @@ def load(self, kern):
         if self._field.is_operator:
             # We always compute operators redundantly out to the L1
             # halo
-            self.set_upper_bound("cell_halo", index=1)
+            self.set_upper_bound("cell_halo", halo_depth=1)
             return
         if (self.field_space.orig_name in
                 const.VALID_DISCONTINUOUS_NAMES):
@@ -275,7 +275,7 @@ def load(self, kern):
             # is therefore no need to iterate into the L1 halo in order
             # to get correct values for annexed dofs.
             if not kern.all_updates_are_writes:
-                self.set_upper_bound("cell_halo", index=1)
+                self.set_upper_bound("cell_halo", halo_depth=1)
                 return
             self.set_upper_bound("ncells")
             return
@@ -286,7 +286,7 @@ def load(self, kern):
             # it is. Again, if the only arguments that are updated have
             # 'GH_WRITE' access then we can relax this condition.
             if not kern.all_updates_are_writes:
-                self.set_upper_bound("cell_halo", index=1)
+                self.set_upper_bound("cell_halo", halo_depth=1)
                 return
             self.set_upper_bound("ncells")
             return
@@ -309,17 +309,18 @@ def set_lower_bound(self, name, index=None):
         self._lower_bound_name = name
         self._lower_bound_index = index
 
-    def set_upper_bound(self, name, index=None):
+    def set_upper_bound(self, name, halo_depth=None):
         '''Set the upper bound of this loop.
 
-        :param name: A loop upper bound name. This should be a supported name.
-        :type name: str
-        :param index: An optional argument indicating the depth of halo
-        :type index: Optional[:py:class:`psyclone.psyir.nodes.Node` | int]
+        :param str name: Loop upper-bound name. Must be a supported name.
+        :param halo_depth: An optional argument indicating the depth of halo
+                           that this loop accesses.
+        :type halo_depth: Optional[:py:class:`psyclone.psyir.nodes.Node` | int]
 
         :raises GenerationError: if supplied with an invalid upper-bound name.
         :raises GenerationError: if supplied with a halo depth < 1.
-
+        :raises TypeError: if the supplied index value is neither an int or
+                           DataNode.
         '''
         const = LFRicConstants()
         if name not in const.VALID_LOOP_BOUNDS_NAMES:
@@ -333,19 +334,23 @@ def set_upper_bound(self, name, index=None):
         # types of bounds, but checking the type of bound as well is a
         # safer option.
         if (name in (["inner"] + const.HALO_ACCESS_LOOP_BOUNDS) and
-                isinstance(index, int)):
-            if index < 1:
+                isinstance(halo_depth, int)):
+            if halo_depth < 1:
                 raise GenerationError(
-                    f"The specified index '{index}' for this upper loop bound "
-                    f"is invalid")
+                    f"The specified halo depth '{halo_depth}' for this loop "
+                    f"upper bound is < 1 which is invalid.")
         self._upper_bound_name = name
-        if index and isinstance(index, int):
+        if halo_depth and isinstance(halo_depth, int):
             # We support specifying depth as an int as a convenience but we
             # now convert it to a PSyIR literal.
-            psyir = Literal(f"{index}", INTEGER_TYPE)
+            psyir = Literal(f"{halo_depth}", INTEGER_TYPE)
             self._upper_bound_halo_depth = psyir
         else:
-            self._upper_bound_halo_depth = index
+            if halo_depth and not isinstance(halo_depth, DataNode):
+                raise TypeError(f"When setting the upper bound of a loop, any "
+                                f"halo depth must be supplied as an int or "
+                                f"PSyIR DataNode but got {type(halo_depth)}")
+            self._upper_bound_halo_depth = halo_depth
 
     @property
     def upper_bound_name(self):

src/psyclone/tests/domain/lfric/lfric_loop_test.py

@@ -128,9 +128,14 @@ def test_set_upper_bound_functions(monkeypatch):
         my_loop.set_upper_bound("start")
     assert "'start' is not a valid upper bound" in str(excinfo.value)
     with pytest.raises(GenerationError) as excinfo:
-        my_loop.set_upper_bound("inner", index=0)
-    assert "specified index" in str(excinfo.value)
-    assert "upper loop bound is invalid" in str(excinfo.value)
+        my_loop.set_upper_bound("inner", halo_depth=0)
+    assert ("specified halo depth '0' for this loop upper bound is < 1 which "
+            "is invalid" in str(excinfo.value))
+    with pytest.raises(TypeError) as excinfo:
+        my_loop.set_upper_bound("inner", halo_depth="wrong")
+    assert ("When setting the upper bound of a loop, any halo depth must be "
+            "supplied as an int or PSyIR DataNode but got "
+            in str(excinfo.value))
 
 
 def test_lower_bound_fortran_1():
@@ -303,7 +308,7 @@ def test_upper_bound_fortran_1():
     psy = PSyFactory(TEST_API, distributed_memory=False).create(invoke_info)
     my_loop = psy.invokes.invoke_list[0].schedule.children[0]
     for option in ["cell_halo", "dof_halo", "inner"]:
-        my_loop.set_upper_bound(option, index=1)
+        my_loop.set_upper_bound(option, halo_depth=1)
         with pytest.raises(GenerationError) as excinfo:
             _ = my_loop._upper_bound_fortran()
             assert (

src/psyclone/tests/dynamo0p3_lma_test.py

@@ -736,7 +736,7 @@ def test_operator_read_level1_halo(tmpdir):
     loop = schedule.children[0]
     # Modify the loop bound so that we attempt to read from the L2 halo
     # (of the operator)
-    loop.set_upper_bound("cell_halo", index=2)
+    loop.set_upper_bound("cell_halo", halo_depth=2)
     # Attempt to generate the code
     with pytest.raises(VisitorError) as excinfo:
         _ = psy.gen