diff --git a/pystencils_tests/test_transformations.py b/pystencils_tests/test_transformations.py index fcd430348ac2661c7fde9c10b55f282ae1caa0cb..ba660a115d6438da31b7bf653730c751da9920a7 100644 --- a/pystencils_tests/test_transformations.py +++ b/pystencils_tests/test_transformations.py @@ -10,6 +10,8 @@ from pystencils.transformations import ( iterate_loops_by_depth, split_inner_loop, loop_blocking ) +from pystencils.cpu import add_pragmas + def test_loop_information(): f, g = ps.fields("f, g: double[2D]") update_rule = ps.Assignment(g[0, 0], f[0, 0]) @@ -115,3 +117,31 @@ def test_split_optimisation(): assert code.count("for") == 6 print(code) + +def test_pragmas(): + f, g = ps.fields("f, g: double[3D]", layout="fzyx") + x = ps.TypedSymbol('x', np.float64) + subs = [ps.Assignment(x, f[0, 0, 0])] + mains = [ps.Assignment(g[0, 0, 0], x)] + ac = ps.AssignmentCollection(mains, subexpressions=subs) + + def prepend_omp_pragmas(ast): + add_pragmas(ast, ["#pragma omp for schedule(dynamic)"], nesting_depth=0) + add_pragmas(ast, ["#pragma omp simd simdlen(8)"], nesting_depth=-1) + + ast_passes = [prepend_omp_pragmas] + + config = ps.CreateKernelConfig(target=ps.Target.CPU, cpu_prepend_optimizations=ast_passes) + ast = ps.create_kernel(ac, config=config) + code = ps.get_code_str(ast) + + assert code.find("#pragma omp for schedule(dynamic)") != -1 + assert code.find("#pragma omp simd simdlen(8)") != -1 + + loops = [loop for loop in filtered_tree_iteration(ast, LoopOverCoordinate, stop_type=SympyAssignment)] + + innermost = list(filter(lambda n: n.is_innermost_loop, loops)) + assert innermost[0].prefix_lines == ["#pragma omp simd simdlen(8)"] + + outermost = list(filter(lambda n: n.is_outermost_loop, loops)) + assert outermost[0].prefix_lines == ["#pragma omp for schedule(dynamic)"]