diff --git a/pystencils_tests/test_opencl.py b/pystencils_tests/test_opencl.py index dfeab790a1129852cc3ac60bdff72cba04015964..5aa6185a6c18126b792d21a5ce3d6f99ff6d4ed4 100644 --- a/pystencils_tests/test_opencl.py +++ b/pystencils_tests/test_opencl.py @@ -39,7 +39,7 @@ def test_print_opencl(): @pytest.mark.skipif(not HAS_OPENCL, reason="Test requires pyopencl") -def test_opencl_jit(): +def test_opencl_jit_fixed_size(): z, y, x = pystencils.fields("z, y, x: [20,30]") assignments = pystencils.AssignmentCollection({ @@ -90,5 +90,57 @@ def test_opencl_jit(): assert np.allclose(result_cuda, result_opencl) +@pytest.mark.skipif(not HAS_OPENCL, reason="Test requires pyopencl") +def test_opencl_jit(): + z, y, x = pystencils.fields("z, y, x: [2d]") + + assignments = pystencils.AssignmentCollection({ + z[0, 0]: x[0, 0] * sp.log(x[0, 0] * y[0, 0]) + }) + + print(assignments) + + ast = pystencils.create_kernel(assignments, target='gpu') + + print(ast) + + code = pystencils.show_code(ast, custom_backend=CudaBackend()) + print(code) + opencl_code = pystencils.show_code(ast, custom_backend=OpenClBackend()) + print(opencl_code) + + cuda_kernel = ast.compile() + assert cuda_kernel is not None + + import pycuda.gpuarray as gpuarray + + x_cpu = np.random.rand(20, 30) + y_cpu = np.random.rand(20, 30) + z_cpu = np.random.rand(20, 30) + + x = gpuarray.to_gpu(x_cpu) + y = gpuarray.to_gpu(y_cpu) + z = gpuarray.to_gpu(z_cpu) + cuda_kernel(x=x, y=y, z=z) + + result_cuda = z.get() + + import pyopencl.array as array + ctx = cl.create_some_context(0) + queue = cl.CommandQueue(ctx) + + x = array.to_device(queue, x_cpu) + y = array.to_device(queue, y_cpu) + z = array.to_device(queue, z_cpu) + + opencl_kernel = make_python_function(ast, queue, ctx) + assert opencl_kernel is not None + opencl_kernel(x=x, y=y, z=z) + + result_opencl = z.get(queue) + + assert np.allclose(result_cuda, result_opencl) + + if __name__ == '__main__': test_opencl_jit()