From ba044dbf3a892ef7b427635c6347abca1dc6c707 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Thu, 16 May 2024 17:32:11 +0200
Subject: [PATCH] add OpenMP insertion to optimize_cpu
---
.../backend/kernelcreation/cpu_optimization.py | 12 +++++++-----
src/pystencils/config.py | 6 +++---
2 files changed, 10 insertions(+), 8 deletions(-)
diff --git a/src/pystencils/backend/kernelcreation/cpu_optimization.py b/src/pystencils/backend/kernelcreation/cpu_optimization.py
index 0e3308cfc..d8d761aec 100644
--- a/src/pystencils/backend/kernelcreation/cpu_optimization.py
+++ b/src/pystencils/backend/kernelcreation/cpu_optimization.py
@@ -1,12 +1,11 @@
from __future__ import annotations
-from typing import cast, TYPE_CHECKING
+from typing import cast
from .context import KernelCreationContext
from ..platforms import GenericCpu
from ..ast.structural import PsBlock
-if TYPE_CHECKING:
- from ...config import CpuOptimConfig
+from ...config import CpuOptimConfig, OpenMpParams
def optimize_cpu(
@@ -33,8 +32,11 @@ def optimize_cpu(
if cfg.vectorize is not False:
raise NotImplementedError("Vectorization not implemented yet")
- if cfg.openmp:
- pass
+ if cfg.openmp is not False:
+ from ..transformations import AddOpenMP
+ params = cfg.openmp if isinstance(cfg.openmp, OpenMpParams) else OpenMpParams()
+ add_omp = AddOpenMP(ctx, params)
+ kernel_ast = cast(PsBlock, add_omp(kernel_ast))
if cfg.use_cacheline_zeroing:
raise NotImplementedError("CL-zeroing not implemented yet")
diff --git a/src/pystencils/config.py b/src/pystencils/config.py
index eaf5e635c..0fc71d666 100644
--- a/src/pystencils/config.py
+++ b/src/pystencils/config.py
@@ -58,11 +58,11 @@ class CpuOptimConfig:
in `CreateKernelConfig.target`, an error will be raised.
"""
- openmp: bool = False
+ openmp: bool | OpenMpParams = False
"""Enable OpenMP parallelization.
- If set to `True`, the kernel will be parallelized using OpenMP according to the OpenMP settings
- given in this configuration.
+ If set to `True`, the kernel will be parallelized using OpenMP according to the default settings in `OpenMpParams`.
+ To customize OpenMP parallelization, pass an instance of `OpenMpParams` instead.
"""
vectorize: bool | VectorizationConfig = False
--
GitLab