#174 add call to set kernel args before each kernel call

stfc · Jun 15, 2018 · b540c75 · b540c75
1 parent e71012f
commit b540c75
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 5 deletions.
diff --git a/src/psyclone/gocean1p0.py b/src/psyclone/gocean1p0.py
@@ -710,7 +710,7 @@ def gen_code(self, parent):
             kernel instance. '''
         from psyclone.f2pygen import CallGen, UseGen
 
-        if self._opencl:
+        if self.root.opencl:
             # OpenCL is completely different so has its own gen method.
             # TODO is there a nicer way of doing this?
             self.gen_ocl(parent)
@@ -756,13 +756,30 @@ def gen_ocl(self, parent):
         '''
         :param parent:
         '''
+        from psyclone.f2pygen import CallGen, DeclGen, AssignGen, CommentGen
+
+        garg = self._find_grid_access()
+        parent.add(DeclGen(parent, datatype="integer",
+                           entity_decls=["globalsize(2)"]))
+        parent.add(AssignGen(
+            parent, lhs="globalsize",
+            rhs="(/{0}%grid%nx, {0}%grid%ny/)".format(garg.name)))
+
+        kernel = "kernel_" + self._name  # TODO use namespace manager
+
+        # First we have to set the kernel arguments
+        parent.add(CallGen(
+            parent, "{0}_set_args".format(self.name),
+            [kernel] + [arg.name for arg in self._arguments.args]))
+        # Then we call clEnqueueNDRangeKernel
         cnull = "C_NULL_PTR"
         cmd_queue = "cmd_queues(1)"  # TODO use namespace manager
-        kernel = "kernel_" + self._name  # TODO use namespace manager
+
         gsize = "globalsize" # TODO use namespace manager
         args = [cmd_queue, kernel, "2", cnull, "C_LOC({0})".format(gsize),
                 cnull, "0", cnull, cnull]
         parent.add(CallGen(parent, "clEnqueueNDRangeKernel", args))
+        parent.add(CommentGen(parent, ""))
 
     @property
     def index_offset(self):

diff --git a/src/psyclone/tests/gocean1p0_opencl_test.py b/src/psyclone/tests/gocean1p0_opencl_test.py
@@ -63,7 +63,31 @@ def test_use_stmts():
     generated_code = str(psy.gen)
     print(generated_code)
     expected = '''\
-  MODULE psy_single_invoke_test
-    USE clfortran
-    USE iso_c_binding'''
+    SUBROUTINE invoke_0_compute_cu(cu_fld, p_fld, u_fld)
+      USE compute_cu_mod, ONLY: compute_cu_code
+      USE clfortran
+      USE iso_c_binding'''
     assert expected in generated_code
+
+
+def test_set_kern_args():
+    ''' Check that we generate the necessary code to set kernel arguments '''
+    _, invoke_info = parse(os.path.join(os.path.
+                                        dirname(os.path.
+                                                abspath(__file__)),
+                                        "test_files", "gocean1p0",
+                                        "single_invoke_two_kernels.f90"),
+                           api=API)
+    psy = PSyFactory(API).create(invoke_info)
+    sched = psy.invokes.invoke_list[0].schedule
+    from psyclone.transformations import OCLTrans
+    otrans = OCLTrans()
+    otrans.apply(sched)
+    generated_code = str(psy.gen)
+    print(generated_code)
+    assert generated_code.count("SUBROUTINE compute_cu_code_set_args(cu_fld, "
+                                "p_fld, u_fld)") == 1
+    assert generated_code.count("SUBROUTINE time_smooth_code_set_args(u_fld, "
+                                "unew_fld, uold_fld)") == 1
+    assert ("CALL compute_cu_code_set_args(cu_fld%data, p_fld%data, "
+            "u_fld%data)" in generated_code)