From a7475e135ddc4b1d8aabf2960b083833978f64c1 Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <kaushikcfd@gmail.com>
Date: Wed, 13 Jul 2022 14:08:36 -0500
Subject: [PATCH] adds test_opencl_offloading

---
 test/unit/test_opencl_offloading.py | 119 ++++++++++++++++++++++++++++
 1 file changed, 119 insertions(+)
 create mode 100644 test/unit/test_opencl_offloading.py

diff --git a/test/unit/test_opencl_offloading.py b/test/unit/test_opencl_offloading.py
new file mode 100644
index 000000000..b39f11a1c
--- /dev/null
+++ b/test/unit/test_opencl_offloading.py
@@ -0,0 +1,119 @@
+import pytest
+pytest.importorskip("pyopencl")
+
+import sys
+import petsc4py
+petsc4py.init(sys.argv
+              + "-viennacl_backend opencl".split()
+              + "-viennacl_opencl_device_type cpu".split())
+from pyop2 import op2
+import pyopencl.array as cla
+import numpy as np
+import loopy as lp
+lp.set_caching_enabled(False)
+
+
+def allclose(a, b, rtol=1e-05, atol=1e-08):
+    """
+    Prefer this routine over np.allclose(...) to allow pycuda/pyopencl arrays
+    """
+    return bool(abs(a - b) < (atol + rtol * abs(b)))
+
+
+def pytest_generate_tests(metafunc):
+    if "backend" in metafunc.fixturenames:
+        from pyop2.backends.opencl import opencl_backend
+        metafunc.parametrize("backend", [opencl_backend])
+
+
+def test_new_backend_raises_not_implemented_error():
+    from pyop2.backends import AbstractComputeBackend
+    unimplemented_backend = AbstractComputeBackend()
+
+    attrs = ["GlobalKernel", "Parloop", "Set", "ExtrudedSet", "MixedSet",
+             "Subset", "DataSet", "MixedDataSet", "Map", "MixedMap", "Dat",
+             "MixedDat", "DatView", "Mat", "Global", "GlobalDataSet",
+             "PETScVecType"]
+
+    for attr in attrs:
+        with pytest.raises(NotImplementedError):
+            getattr(unimplemented_backend, attr)
+
+
+def test_dat_with_petscvec_representation(backend):
+    op2.set_offloading_backend(backend)
+
+    nelems = 9
+    data = np.random.rand(nelems)
+    set_ = op2.compute_backend.Set(nelems)
+    dset = op2.compute_backend.DataSet(set_, 1)
+    dat = op2.compute_backend.Dat(dset, data.copy())
+
+    assert isinstance(dat.data_ro, np.ndarray)
+    dat.data[:] *= 3
+
+    with op2.offloading():
+        assert isinstance(dat.data_ro, cla.Array)
+        dat.data[:] *= 2
+
+    assert isinstance(dat.data_ro, np.ndarray)
+    np.testing.assert_allclose(dat.data_ro, 6*data)
+
+
+def test_dat_not_as_petscvec(backend):
+    op2.set_offloading_backend(backend)
+
+    nelems = 9
+    data = np.random.randint(low=-10, high=10,
+                             size=nelems,
+                             dtype=np.int64)
+    set_ = op2.compute_backend.Set(nelems)
+    dset = op2.compute_backend.DataSet(set_, 1)
+    dat = op2.compute_backend.Dat(dset, data.copy())
+
+    assert isinstance(dat.data_ro, np.ndarray)
+    dat.data[:] *= 3
+
+    with op2.offloading():
+        assert isinstance(dat.data_ro, cla.Array)
+        dat.data[:] *= 2
+
+    assert isinstance(dat.data_ro, np.ndarray)
+    np.testing.assert_allclose(dat.data_ro, 6*data)
+
+
+def test_global_reductions(backend):
+    op2.set_offloading_backend(backend)
+
+    sum_knl = lp.make_function(
+        "{ : }",
+        """
+        g[0] = g[0] + x[0]
+        """,
+        [lp.GlobalArg("g,x",
+                      dtype="float64",
+                      shape=lp.auto,
+                      )],
+        name="sum_knl",
+        target=lp.CWithGNULibcTarget(),
+        lang_version=(2018, 2))
+
+    rng = np.random.default_rng()
+    nelems = 4_000
+    data_to_sum = rng.random(nelems)
+
+    with op2.offloading():
+
+        set_ = op2.compute_backend.Set(4_000)
+        dset = op2.compute_backend.DataSet(set_, 1)
+
+        dat = op2.compute_backend.Dat(dset, data_to_sum.copy())
+        glob = op2.compute_backend.Global(1, 0, np.float64, "g")
+
+        op2.parloop(op2.Kernel(sum_knl, "sum_knl"),
+                    set_,
+                    glob(op2.INC),
+                    dat(op2.READ))
+
+        assert isinstance(glob.data_ro, cla.Array)
+        assert allclose(glob.data_ro[0], data_to_sum.sum())