diff --git a/tests/kernelcreation/test_reduction.py b/tests/kernelcreation/test_reduction.py index 537eb4b67ade9f38fdaa05626c6dc29511ad2cad..9fd385e02beb0a730cb7f5702a494fc50ab867e3 100644 --- a/tests/kernelcreation/test_reduction.py +++ b/tests/kernelcreation/test_reduction.py @@ -17,38 +17,50 @@ SOLUTION = { } -@pytest.mark.parametrize('dtype', ["float64"]) -@pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"]) -def test_reduction(target, dtype, op): - gpu_avail = target is ps.Target.GPU - +# get AST for kernel with reduction assignment +def get_reduction_assign_ast(dtype, op, config): x = ps.fields(f'x: {dtype}[1d]') w = ps.TypedSymbol("w", dtype) - # kernel with reduction assignment - red_assign = reduction_assignment_from_str(w, op, x.center()) - vectorize_info = {'instruction_set': 'avx', 'assume_inner_stride_one': True} + return ps.create_kernel([red_assign], config, default_dtype=dtype) + - config = ps.CreateKernelConfig(target=target) if gpu_avail \ - else ps.CreateKernelConfig(target=target, cpu_openmp=True, cpu_vectorize_info=vectorize_info) +@pytest.mark.parametrize('instruction_set', ['sse', 'avx']) +@pytest.mark.parametrize('dtype', ["float64", "float32"]) +@pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"]) +def test_reduction_cpu(instruction_set, dtype, op): + + vectorize_info = {'instruction_set': instruction_set, 'assume_inner_stride_one': True} - ast_reduction = ps.create_kernel([red_assign], config, default_dtype=dtype) + config = ps.CreateKernelConfig(target=ps.Target.CPU, cpu_openmp=True, cpu_vectorize_info=vectorize_info) + + ast_reduction = get_reduction_assign_ast(dtype, op, config) ps.show_code(ast_reduction) + kernel_reduction = ast_reduction.compile() - # code_reduction = ps.get_code_str(ast_reduction) + array = np.full((SIZE,), INIT_ARR, dtype=dtype) + reduction_array = np.full((1,), INIT_W, dtype=dtype) + + kernel_reduction(x=array, w=reduction_array) + assert np.allclose(reduction_array, SOLUTION[op]) + + +@pytest.mark.parametrize('dtype', ["float64", "float32"]) +@pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"]) +def test_reduction_gpu(dtype, op): + config = ps.CreateKernelConfig(target=ps.Target.GPU) + + ast_reduction = get_reduction_assign_ast(dtype, op, config) + ps.show_code(ast_reduction) kernel_reduction = ast_reduction.compile() array = np.full((SIZE,), INIT_ARR, dtype=dtype) reduction_array = np.full((1,), INIT_W, dtype=dtype) - if gpu_avail: - array_gpu = cp.asarray(array) - reduction_array_gpu = cp.asarray(reduction_array) + 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]) + kernel_reduction(x=array_gpu, w=reduction_array_gpu) + assert np.allclose(reduction_array_gpu.get(), SOLUTION[op])