Split reduction test into separate CPU/GPU tests

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])