Add reductions param to AddOpenMP

src/pystencils/backend/transformations/add_pragmas.py

@@ -8,6 +8,7 @@ from ..kernelcreation import KernelCreationContext
 from ..ast import PsAstNode
 from ..ast.structural import PsBlock, PsLoop, PsPragma, PsStructuralNode
 from ..ast.expressions import PsExpression
+from ..kernelcreation.context import ReductionInfo
 
 from ...types import PsScalarType
 
@@ -104,12 +105,16 @@ class AddOpenMP:
     def __init__(
         self,
         ctx: KernelCreationContext,
+        reductions: Sequence[ReductionInfo] = (),
         nesting_depth: int = 0,
         num_threads: int | None = None,
         schedule: str | None = None,
         collapse: int | None = None,
         omit_parallel: bool = False,
     ) -> None:
+
+        self._reductions = reductions
+
         pragma_text = "omp"
 
         if not omit_parallel:
@@ -123,16 +128,15 @@ class AddOpenMP:
         if num_threads is not None:
             pragma_text += f" num_threads({str(num_threads)})"
 
-        if bool(ctx.reduction_data):
-            for _, reduction_info in ctx.reduction_data.items():
-                if isinstance(reduction_info.local_symbol.dtype, PsScalarType):
-                    pragma_text += (
-                        f" reduction({reduction_info.op.value}: {reduction_info.local_symbol.name})"
-                    )
-                else:
-                    NotImplementedError(
-                        "OMP: Reductions for non-scalar data types are not supported yet."
-                    )
+        for reduction_info in self._reductions:
+            if isinstance(reduction_info.local_symbol.dtype, PsScalarType):
+                pragma_text += (
+                    f" reduction({reduction_info.op.value}: {reduction_info.local_symbol.name})"
+                )
+            else:
+                NotImplementedError(
+                    "OMP: Reductions for non-scalar data types are not supported yet."
+                )
 
         if collapse is not None:
             if collapse <= 0:

src/pystencils/codegen/driver.py

@@ -359,6 +359,7 @@ class DefaultKernelCreationDriver:
 
             add_omp = AddOpenMP(
                 self._ctx,
+                reductions=list(self._ctx.reduction_data.values()),
                 nesting_depth=omp_options.get_option("nesting_depth"),
                 num_threads=omp_options.get_option("num_threads"),
                 schedule=omp_options.get_option("schedule"),