diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index ded36b50aa972c808da3908c5920c20af9b78208..4122e8e73b275206cb9f87e9031a4b26a6be2b14 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -10,7 +10,7 @@ hyteg-hfg-integration-test: - GIT_CONFIG_COUNT=1 GIT_CONFIG_KEY_0="url.https://gitlab-ci-token:${CI_JOB_TOKEN}@i10git.cs.fau.de.insteadOf" GIT_CONFIG_VALUE_0="ssh://git@i10git.cs.fau.de" python -m pip install -r requirements.txt - pip freeze - rm -rf ../operators/* - - python3 generate.py -o ../operators ../operators.toml + - python3 generate.py -o ../operators ../operators.toml --processes 8 - cd .. - git status - git diff diff --git a/generate/generate.py b/generate/generate.py index 308e6b38996893c3fe6e1427f55ea06616a99dec..ad89df2165790c622ac46d8a25a2b912ff35fa23 100644 --- a/generate/generate.py +++ b/generate/generate.py @@ -3,7 +3,7 @@ from functools import partial import os import sys from typing import Any, Dict, List, Union - +from multiprocessing import Pool if sys.version_info >= (3, 11): import tomllib else: @@ -37,6 +37,14 @@ def parse_args() -> argparse.Namespace: help="Generates cmake files from existing source and header files and quits.", ) + parser.add_argument( + "--processes", + action="store", + type=int, + default=1, + help="Generates kernels in parallel with the specified number of processes.", + ) + return parser.parse_args() @@ -81,19 +89,16 @@ def main() -> None: os.makedirs(args.output, exist_ok=True) if not args.cmake_only: - with open(args.filename, "rb") as f: + with (open(args.filename, "rb") as f): toml_dict = tomllib.load(f) toml_dict = unfold_toml_dict(toml_dict) - for form_str, operators in toml_dict.items(): - kernel_implementations = {} - for spec in operators: - op_kernel_impls = generate_operator(args, form_str, spec) - - for platform, impls in op_kernel_impls.items(): - if not platform in kernel_implementations: - kernel_implementations[platform] = [] - kernel_implementations[platform].extend(impls) + with Pool(args.processes) as pool: + for form_str, operators in toml_dict.items(): + for spec in operators: + pool.apply_async(generate_operator, (args, form_str, spec)) + pool.close() + pool.join() generate_cmake_from_cpp_files(args.output)