import numpy as np import numpy.linalg as la import pyopencl as cl import pyopencl.array # noqa import pyopencl.clrandom # noqa import loopy as lp import sys from pyopencl.tools import ( # noqa pytest_generate_tests_for_pyopencl as pytest_generate_tests) _WENO_PRG = [] queue = [] def get_weno_program(): if _WENO_PRG: return _WENO_PRG[0] fn = "WENO.F90" with open(fn, "r") as infile: infile_content = infile.read() prg = lp.parse_transformed_fortran(infile_content, filename=fn) _WENO_PRG.append(prg) return prg def with_root_kernel(prg, root_name): # FIXME This is a little less beautiful than it could be new_prg = prg.copy(name=root_name) for name in prg: clbl = new_prg[name] if isinstance(clbl, lp.LoopKernel) and clbl.is_called_from_host: new_prg = new_prg.with_kernel(clbl.copy(is_called_from_host=False)) new_prg = new_prg.with_kernel(prg[root_name].copy(is_called_from_host=True)) return new_prg def test_setup(ctx_factory): """hack so we only need to set the context one time -- BUT, can we guarantee this runs before all other tests? """ ctx = ctx_factory() queue.append(cl.CommandQueue(ctx)) def test_matvec(): #def test_matvec(ctx_factory): #ctx = ctx_factory() #queue = cl.CommandQueue(ctx) prg = get_weno_program() prg = with_root_kernel(prg, "mult_mat_vec") a = np.random.rand(10, 10).astype(np.float32).copy(order="F") b = np.random.rand(10).astype(np.float32) c_dev = cl.array.empty(queue[0], 10, dtype=np.float32) prg = lp.set_options(prg, write_cl=True) prg(queue[0], a=a, b=b, c=c_dev, alpha=1.0) c = c_dev.get() assert la.norm(a@b - c, 2)/la.norm(c) < 1e-5 def test_compute_flux_derivatives(): #def test_compute_flux_derivatives(ctx_factory): #ctx = ctx_factory() #queue = cl.CommandQueue(ctx) ndim = 3 nvars = 5 nx = 10 ny = 10 nz = 10 prg = get_weno_program() prg = lp.fix_parameters(prg, nx=nx, ny=ny, nz=nz, _remove=False) states = np.random.randn(nvars, nx+6, ny+6, nz+6).astype(np.float32).copy( order="F") fluxes = np.random.randn(nvars, ndim, nx+6, ny+6, nz+6).astype(np.float32).copy( order="F") metrics = np.random.randn(ndim, ndim, nx+6, ny+6, nz+6).astype(np.float32).copy( order="F") metric_jacobians = np.random.randn(nx+6, ny+6, nz+6).astype(np.float32).copy( order="F") flux_derivatives_dev = cl.array.empty(queue[0], (nvars, ndim, nx+6, ny+6, nz+6), dtype=np.float32, order="F") prg(queue[0], nvars=nvars, ndim=ndim, nx=nx, ny=ny, nz=nz, states=states, fluxes=fluxes, metrics=metrics, metric_jacobians=metric_jacobians, flux_derivatives=flux_derivatives_dev) print('Ran with some dummy arguments') # This lets you run 'python test.py test_case(cl._csc)' without pytest. if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1]) else: from pytest import main main([__file__])