Add function to prepend pragmas to loops

Add a function to prepend #pragma directives to domain loops according to their nesting depth. Useful for compiler-specific annotations and interactions, especially with less well-known niche compilers (like Fujitsu's ARM-compiler).

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