Merge pull request #520 from OP2/mesh-level-wrapper

Allow passing layer argument to kernel

Author: miklos1

pyop2/base.py

@@ -4040,6 +4040,14 @@ def __init__(self, kernel, iterset, *args, **kwargs):
         self._kernel = kernel
         self._is_layered = iterset._extruded
         self._iteration_region = kwargs.get("iterate", None)
+        self._pass_layer_arg = kwargs.get("pass_layer_arg", False)
+
+        if self._pass_layer_arg:
+            if self.is_direct:
+                raise ValueError("Can't request layer arg for direct iteration")
+            if not self._is_layered:
+                raise ValueError("Can't request layer arg for non-extruded iteration")
+
         # Are we only computing over owned set entities?
         self._only_local = isinstance(iterset, LocalSet)
 
@@ -4389,6 +4397,10 @@ def par_loop(kernel, it_space, *args, **kwargs):
                  except the top layer, accessing data two adjacent (in
                  the extruded direction) cells at a time.
 
+    :kwarg pass_layer_arg: Should the wrapper pass the current layer
+        into the kernel (as an ``int``). Only makes sense for
+        indirect extruded iteration.
+
     .. warning ::
         It is the caller's responsibility that the number and type of all
         :class:`base.Arg`\s passed to the :func:`par_loop` match those expected

pyop2/sequential.py

@@ -609,6 +609,7 @@ def __init__(self, kernel, itspace, *args, **kwargs):
         self._args = args
         self._direct = kwargs.get('direct', False)
         self._iteration_region = kwargs.get('iterate', ALL)
+        self._pass_layer_arg = kwargs.get('pass_layer_arg', False)
         # Copy the class variables, so we don't overwrite them
         self._cppargs = dcopy(type(self)._cppargs)
         self._libraries = dcopy(type(self)._libraries)
@@ -709,7 +710,8 @@ def generate_code(self):
                                                kernel_name=self._kernel._name,
                                                user_code=self._kernel._user_code,
                                                wrapper_name=self._wrapper_name,
-                                               iteration_region=self._iteration_region)
+                                               iteration_region=self._iteration_region,
+                                               pass_layer_arg=self._pass_layer_arg)
         return self._code_dict
 
     def set_argtypes(self, iterset, *args):
@@ -781,7 +783,8 @@ def prepare_arglist(self, iterset, *args):
     @cached_property
     def _jitmodule(self):
         return JITModule(self.kernel, self.it_space, *self.args,
-                         direct=self.is_direct, iterate=self.iteration_region)
+                         direct=self.is_direct, iterate=self.iteration_region,
+                         pass_layer_arg=self._pass_layer_arg)
 
     @collective
     def _compute(self, part, fun, *arglist):
@@ -792,7 +795,7 @@ def _compute(self, part, fun, *arglist):
 
 def wrapper_snippets(itspace, args,
                      kernel_name=None, wrapper_name=None, user_code=None,
-                     iteration_region=ALL):
+                     iteration_region=ALL, pass_layer_arg=False):
     """Generates code snippets for the wrapper,
     ready to be into a template.
 
@@ -914,6 +917,10 @@ def extrusion_loop():
             _buf_gather[arg] = "\n".join([_itspace_loops, _buf_gather[arg], _itspace_loop_close])
     _kernel_args = ', '.join([arg.c_kernel_arg(count) if not arg._uses_itspace else _buf_name[arg]
                               for count, arg in enumerate(args)])
+
+    if pass_layer_arg:
+        _kernel_args += ", j_0"
+
     _buf_gather = ";\n".join(_buf_gather.values())
     _buf_decl = ";\n".join(_buf_decl.values())
 

test/unit/test_extrusion.py

@@ -377,6 +377,23 @@ def test_direct_loop_inc(self, xtr_nodes):
                      dat.dataset.set, dat(op2.INC))
         assert numpy.allclose(dat.data[:], 1.0)
 
+    def test_extruded_layer_arg(self, elements, field_map, dat_f):
+        """Tests that the layer argument is being passed when prompted
+        to in the parloop."""
+
+        kernel_blah = """void kernel_blah(double* x[], int layer_arg){
+                                                 x[0][0] = layer_arg;
+                                              }\n"""
+
+        op2.par_loop(op2.Kernel(kernel_blah, "kernel_blah"),
+                     elements, dat_f(op2.WRITE, field_map),
+                     pass_layer_arg=True)
+        end = layers - 1
+        start = 0
+        ref = np.arange(start, end)
+        assert [dat_f.data[end*n:end*(n+1)] == ref
+                for n in range(len(dat_f.data) + 1)]
+
     def test_write_data_field(self, elements, dat_coords, dat_field, coords_map, field_map, dat_f):
         kernel_wo = "void kernel_wo(double* x[]) { x[0][0] = 42.0; }\n"