compiler: Patch forgotten openmp-offloading handler

devitocodes · Jan 31, 2022 · d8e5d4a · d8e5d4a
1 parent 67b6cff
commit d8e5d4a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/devito/passes/iet/languages/openmp.py b/devito/passes/iet/languages/openmp.py
@@ -159,8 +159,8 @@ class OmpBB(PragmaLangBB):
                                        DefFunction('omp_get_initial_device')]),
         'device-get':
             Call('omp_get_default_device'),
-        'device-alloc': lambda i, j, *a, **kw:
-            Call('omp_target_alloc', (i, j)),
+        'device-alloc': lambda i, j, retobj:
+            Call('omp_target_alloc', (i, j), retobj=retobj, cast=True),
         'device-free': lambda i, j:
             Call('omp_target_free', (i, j))
     }

diff --git a/tests/test_gpu_openmp.py b/tests/test_gpu_openmp.py
@@ -211,15 +211,15 @@ def test_array_rw(self):
 
         assert len(op.body.allocs) == 1
         assert str(op.body.allocs[0]) ==\
-            ('float *r0_vec = (float*) '
-             'omp_target_alloc(sizeof(float[x_size*y_size*z_size]), '
+            ('float * r0_vec = (float *)'
+             'omp_target_alloc(x_size*y_size*z_size*sizeof(float),'
              'omp_get_default_device());')
         assert len(op.body.maps) == 2
         assert all('r0' not in str(i) for i in op.body.maps)
 
         assert len(op.body.frees) == 1
         assert str(op.body.frees[0]) ==\
-            'omp_target_free(r0_vec, omp_get_default_device());'
+            'omp_target_free(r0_vec,omp_get_default_device());'
         assert len(op.body.unmaps) == 3
         assert all('r0' not in str(i) for i in op.body.unmaps)