-
Notifications
You must be signed in to change notification settings - Fork 46
/
config.py
88 lines (71 loc) · 3.29 KB
/
config.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2023, ETH Zurich
# All rights reserved.
#
# This file is part of the GT4Py project and the GridTools framework.
# GT4Py is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or any later
# version. See the LICENSE.txt file at the top-level directory of this
# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later
import multiprocessing
import os
from typing import Any, Dict, List, Optional
import gridtools_cpp
GT4PY_INSTALLATION_PATH: str = os.path.dirname(os.path.abspath(__file__))
# Default paths (taken from user's environment vars when possible)
BOOST_ROOT: str = os.environ.get(
"BOOST_ROOT", os.environ.get("BOOST_HOME", os.path.abspath("/usr/local"))
)
CUDA_ROOT: str = os.environ.get(
"CUDA_HOME", os.environ.get("CUDA_PATH", os.path.abspath("/usr/local/cuda"))
)
CUDA_HOST_CXX: Optional[str] = os.environ.get("CUDA_HOST_CXX", None)
if "GT4PY_USE_HIP" in os.environ:
GT4PY_USE_HIP: bool = bool(int(os.environ["GT4PY_USE_HIP"]))
else:
# Autodetect cupy with ROCm/HIP support
try:
import cupy as _cp
GT4PY_USE_HIP = _cp.cuda.get_hipcc_path() is not None
del _cp
except Exception:
GT4PY_USE_HIP = False
GT_INCLUDE_PATH: str = os.path.abspath(gridtools_cpp.get_include_dir())
GT_CPP_TEMPLATE_DEPTH: int = 1024
# Settings dict
GT4PY_EXTRA_COMPILE_ARGS: str = os.environ.get("GT4PY_EXTRA_COMPILE_ARGS", "")
extra_compile_args: List[str] = (
list(GT4PY_EXTRA_COMPILE_ARGS.split(" ")) if GT4PY_EXTRA_COMPILE_ARGS else []
)
build_settings: Dict[str, Any] = {
"boost_include_path": os.path.join(BOOST_ROOT, "include"),
"cuda_bin_path": os.path.join(CUDA_ROOT, "bin"),
"cuda_include_path": os.path.join(CUDA_ROOT, "include"),
"cuda_arch": os.environ.get("CUDA_ARCH", None),
"gt_include_path": os.environ.get("GT_INCLUDE_PATH", GT_INCLUDE_PATH),
"openmp_cppflags": os.environ.get("OPENMP_CPPFLAGS", "-fopenmp").split(),
"openmp_ldflags": os.environ.get("OPENMP_LDFLAGS", "-fopenmp").split(),
"extra_compile_args": {"cxx": extra_compile_args, "cuda": extra_compile_args},
"extra_link_args": [],
"parallel_jobs": multiprocessing.cpu_count(),
"cpp_template_depth": os.environ.get("GT_CPP_TEMPLATE_DEPTH", GT_CPP_TEMPLATE_DEPTH),
}
if GT4PY_USE_HIP:
build_settings["cuda_library_path"] = os.path.join(CUDA_ROOT, "lib")
else:
build_settings["cuda_library_path"] = os.path.join(CUDA_ROOT, "lib64")
if CUDA_HOST_CXX is not None:
build_settings["extra_compile_args"]["cuda"].append(f"-ccbin={CUDA_HOST_CXX}")
cache_settings: Dict[str, Any] = {
"dir_name": os.environ.get("GT_CACHE_DIR_NAME", ".gt_cache"),
"root_path": os.environ.get("GT_CACHE_ROOT", os.path.abspath(".")),
"load_retries": int(os.environ.get("GT_CACHE_LOAD_RETRIES", 3)),
"load_retry_delay": int(os.environ.get("GT_CACHE_LOAD_RETRY_DELAY", 100)), # unit milliseconds
}
code_settings: Dict[str, Any] = {"root_package_name": "_GT_"}
os.environ.setdefault("DACE_CONFIG", os.path.join(os.path.abspath("."), ".dace.conf"))
DACE_DEFAULT_BLOCK_SIZE: str = os.environ.get("DACE_DEFAULT_BLOCK_SIZE", "64,8,1")