Skip to content
Snippets Groups Projects
Commit 4c726aa6 authored by Richard Angersbach's avatar Richard Angersbach
Browse files

Prepare reduction test for GPU support

parent b352a2e2
1 merge request!438Reduction Support
Pipeline #72630 failed with stages
in 9 minutes and 49 seconds
This commit is part of merge request !438. Comments created here will be created in the context of that merge request.
import pytest import pytest
import numpy as np import numpy as np
import sympy as sp import sympy as sp
import cupy as cp
import pystencils as ps import pystencils as ps
from pystencils.sympyextensions import reduced_assign from pystencils.sympyextensions import reduced_assign
...@@ -18,6 +19,9 @@ SOLUTION = { ...@@ -18,6 +19,9 @@ SOLUTION = {
@pytest.mark.parametrize('dtype', ["float64"]) @pytest.mark.parametrize('dtype', ["float64"])
@pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"]) @pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"])
def test_reduction(dtype, op): def test_reduction(dtype, op):
gpu_avail = True
x = ps.fields(f'x: {dtype}[1d]') x = ps.fields(f'x: {dtype}[1d]')
w = sp.Symbol("w") w = sp.Symbol("w")
...@@ -25,7 +29,7 @@ def test_reduction(dtype, op): ...@@ -25,7 +29,7 @@ def test_reduction(dtype, op):
reduction_assignment = reduced_assign(w, op, x.center()) reduction_assignment = reduced_assign(w, op, x.center())
config = ps.CreateKernelConfig(cpu_openmp=True) config = ps.CreateKernelConfig(target=ps.Target.GPU) if gpu_avail else ps.CreateKernelConfig(cpu_openmp=True)
ast_reduction = ps.create_kernel([reduction_assignment], config, default_dtype=dtype) ast_reduction = ps.create_kernel([reduction_assignment], config, default_dtype=dtype)
#code_reduction = ps.get_code_str(ast_reduction) #code_reduction = ps.get_code_str(ast_reduction)
...@@ -35,5 +39,13 @@ def test_reduction(dtype, op): ...@@ -35,5 +39,13 @@ def test_reduction(dtype, op):
array = np.full((SIZE,), INIT, dtype=dtype) array = np.full((SIZE,), INIT, dtype=dtype)
reduction_array = np.zeros(1, dtype=dtype) reduction_array = np.zeros(1, dtype=dtype)
kernel_reduction(x=array, w=reduction_array)
assert np.allclose(reduction_array, SOLUTION[op]) if gpu_avail:
\ No newline at end of file array_gpu = cp.asarray(array)
reduction_array_gpu = cp.asarray(reduction_array)
kernel_reduction(x=array_gpu, w=reduction_array_gpu)
assert np.allclose(reduction_array_gpu.get(), SOLUTION[op])
else:
kernel_reduction(x=array, w=reduction_array)
assert np.allclose(reduction_array, SOLUTION[op])
\ No newline at end of file
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment