Skip to content

Commit e0968b7

Browse files
casparvlcasparvl
authored
Merge pull request #228 from zerefwayne/gpu-tag-rocm
Add gpu tag for ROCm-LLVM and refactor method
2 parents 4309853 + 8a0468e commit e0968b7

2 files changed

Lines changed: 132 additions & 23 deletions

File tree

‎.github/workflows/test-eb-hooks.yml‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,3 +116,55 @@ jobs:
116116
eb --hooks=$PWD/eb_hooks.py "$INCOMPATIBLE_EASYCONFIG" --stop fetch 2>&1 1>/dev/null | grep -q "does not contain a valid list of dictionaries"
117117
echo "Incorrect format for EESSI_SITE_TOP_LEVEL_TOOLCHAINS caught"
118118
119+
check_inject_gpu_property:
120+
runs-on: ubuntu-24.04
121+
strategy:
122+
matrix:
123+
# ROCm-LLVM / rocm-compilers / rompi only exist in the 2025.06 stack
124+
EESSI_VERSION:
125+
- '2025.06'
126+
127+
steps:
128+
- name: Check out software-layer repository
129+
uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
130+
131+
- name: Mount EESSI CernVM-FS repository
132+
uses: eessi/github-action-eessi@v3
133+
with:
134+
eessi_stack_version: ${{matrix.EESSI_VERSION}}
135+
use_eessi_module: true
136+
137+
- name: Test that inject_gpu_property tags GPU software correctly
138+
if: ${{ github.event_name == 'pull_request' }}
139+
run: |
140+
module load EESSI-extend/${{matrix.EESSI_VERSION}}-easybuild
141+
142+
echo "Running inject_gpu_property tests for EESSI/${{matrix.EESSI_VERSION}}..."
143+
144+
# Control: a non-GPU easyconfig must not receive the gpu property
145+
eb --hooks=$PWD/eb_hooks.py --extended-dry-run M4-1.4.19-GCCcore-14.2.0.eb 2>&1 | grep -qiF 'add_property("arch","gpu")' && exit 1
146+
echo "Non-GPU easyconfig untouched"
147+
148+
# ROCm-LLVM as a direct dependency (rocm-compilers bundle)
149+
echo "Testing ROCm-LLVM as a direct dependency"
150+
OUT=$(eb --hooks=$PWD/eb_hooks.py --extended-dry-run rocm-compilers-19.0.0-ROCm-6.4.1.eb 2>&1)
151+
echo "$OUT" | grep -qiF 'add_property("arch","gpu")'
152+
echo "$OUT" | grep -qiF 'setenv("EESSIROCMVERSION","6.4.1")'
153+
echo "Direct ROCm-LLVM dependency detected"
154+
155+
# rocm-compilers as the toolchain
156+
echo "Testing ROCm-LLVM in a rocm-compilers toolchain"
157+
OUT=$(eb --hooks=$PWD/eb_hooks.py --extended-dry-run rocBLAS-4.4.0-rocm-compilers-19.0.0-ROCm-6.4.1.eb 2>&1)
158+
echo "$OUT" | grep -qiF 'add_property("arch","gpu")'
159+
echo "$OUT" | grep -qiF 'setenv("EESSIROCMVERSION","6.4.1")'
160+
echo "rocm-compilers toolchain handled"
161+
162+
# TODO: rompi test to be brought back once rompi/2025a is available in EESSI
163+
164+
# # rompi as the toolchain (ROCm-LLVM nested inside rocm-compilers bundle)
165+
# echo "Testing ROCm-LLVM in a rompi toolchain (rocm-compilers bundle)"
166+
# # added --disable-cleanup-tmpdir to ignore irrelevant EasyBuild cleanup error
167+
# OUT=$(eb --hooks=$PWD/eb_hooks.py --extended-dry-run OSU-Micro-Benchmarks-7.5-rompi-2025a.eb 2>&1)
168+
# echo "$OUT" | grep -qiF 'add_property("arch","gpu")'
169+
# echo "$OUT" | grep -qiF 'setenv("EESSIROCMVERSION","6.4.1")'
170+
# echo "rompi toolchain handled"

‎eb_hooks.py‎

Lines changed: 80 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
import easybuild.tools.environment as env
1212
from easybuild.easyblocks.generic.configuremake import obtain_config_guess
1313
from easybuild.framework.easyconfig.constants import EASYCONFIG_CONSTANTS
14-
from easybuild.framework.easyconfig.easyconfig import get_toolchain_hierarchy
14+
from easybuild.framework.easyconfig.easyconfig import (
15+
get_toolchain_hierarchy,
16+
)
1517
from easybuild.tools import config
1618
from easybuild.tools.build_log import EasyBuildError, print_msg, print_warning
1719
from easybuild.tools.config import build_option, install_path, update_build_option
@@ -2014,38 +2016,93 @@ def replace_binary_non_distributable_files_with_symlinks(log, install_dir, pkg_n
20142016
symlink(host_inj_path, full_path)
20152017

20162018

2019+
def find_rocm_llvm_dependency(ec):
2020+
"""
2021+
Return the ROCm-LLVM dependency for this easyconfig, or None. ROCm-LLVM can
2022+
be a direct dependency, a direct toolchain component (rocm-compilers as the
2023+
toolchain), or one level deeper inside the rocm-compilers bundle when the
2024+
toolchain is rompi/rfbf/rfoss.
2025+
"""
2026+
# Check if ROCm-LLVM is a direct dependency, and if so, return that
2027+
#
2028+
for dep in ec.asdict()['dependencies']:
2029+
# dep is a tuple (name, version, versionsuffix, toolchain); normalise it to the
2030+
# same dict format as ec.toolchain.tcdeps entries before returning
2031+
if dep[0] == 'ROCm-LLVM':
2032+
return {
2033+
'name': dep[0],
2034+
'version': dep[1],
2035+
'versionsuffix': dep[2] if len(dep) > 2 else '',
2036+
'toolchain': dep[3] if len(dep) > 3 else None,
2037+
}
2038+
2039+
# ROCm-LLVM can also be part of the toolchain. First, return early if this is NOT a ROCm-based toolchain
2040+
if ec['toolchain']['name'] not in ('rocm-compilers', 'rompi', 'rfbf', 'rfoss'):
2041+
return None
2042+
2043+
tcdeps = ec.toolchain.tcdeps or []
2044+
# Check if ROCm-LLVM is a direct dependency for this toolchain (which would be the case for rocm-compilers)
2045+
for dep in tcdeps:
2046+
if dep['name'] == 'ROCm-LLVM':
2047+
return dep
2048+
# For rompi, rfbf, rfoss, ROCm-LLVM is pulled in indirectly via rocm-compilers. the rocm-compilers
2049+
# toolchain dependency already encodes the ROCm version in its version string (e.g. '19.0.0-ROCm-6.4.1')
2050+
rocm_prefix = '-ROCm-'
2051+
for dep in tcdeps:
2052+
if dep['name'] == 'rocm-compilers':
2053+
full_version = dep['version'] + dep.get('versionsuffix', '')
2054+
if rocm_prefix in full_version:
2055+
version, rocm_version = full_version.split(rocm_prefix, 1)
2056+
return {
2057+
'name': dep['name'],
2058+
'version': version,
2059+
'versionsuffix': rocm_prefix + rocm_version,
2060+
}
2061+
2062+
return None
2063+
2064+
20172065
def inject_gpu_property(ec):
20182066
"""
20192067
Add 'gpu' property and EESSI<PACKAGE>VERSION envvars via modluafooter
20202068
easyconfig parameter, and drop dependencies to build dependencies
20212069
"""
20222070
ec_dict = ec.asdict()
2023-
# Check if CUDA, cuDNN, you-name-it is in the dependencies, if so
2024-
# - drop dependency to build dependency
2025-
# - add 'gpu' Lmod property
2026-
# - add envvar with package version
2027-
pkg_names = ( "CUDA", "cuDNN" )
20282071
pkg_versions = { }
20292072
add_gpu_property = ''
20302073

2031-
for pkg_name in pkg_names:
2032-
# Check if pkg_name is in the dependencies, if so drop dependency to build
2033-
# dependency and set variable for later adding the 'gpu' Lmod property
2034-
# to '.remove' dependencies from ec_dict['dependencies'] we make a copy,
2035-
# iterate over the copy and can then savely use '.remove' on the original
2036-
# ec_dict['dependencies'].
2037-
deps = ec_dict['dependencies'][:]
2038-
if (pkg_name in [dep[0] for dep in deps]):
2074+
# Check if pkg_name is related to CUDA, if so drop dependency to build
2075+
# dependency and set variable for later adding the 'gpu' Lmod property
2076+
# to '.remove' dependencies from ec_dict['dependencies'] we make a copy,
2077+
# iterate over the copy and can then savely use '.remove' on the original
2078+
# ec_dict['dependencies'].
2079+
for pkg_name in ('CUDA', 'cuDNN'):
2080+
for dep in ec_dict['dependencies'][:]:
2081+
if dep[0] != pkg_name:
2082+
continue
2083+
20392084
add_gpu_property = 'add_property("arch","gpu")'
2040-
for dep in deps:
2041-
if pkg_name == dep[0]:
2042-
# make pkg_name a build dependency only (rpathing saves us from link errors)
2043-
ec.log.info("Dropping dependency on %s to build dependency" % pkg_name)
2044-
ec_dict['dependencies'].remove(dep)
2045-
if dep not in ec_dict['builddependencies']:
2046-
ec_dict['builddependencies'].append(dep)
2047-
# take note of version for creating the modluafooter
2048-
pkg_versions[pkg_name] = dep[1]
2085+
pkg_versions[pkg_name] = dep[1]
2086+
2087+
ec.log.info("Dropping dependency on %s to build dependency" % pkg_name)
2088+
ec_dict['dependencies'].remove(dep)
2089+
if dep not in ec_dict['builddependencies']:
2090+
ec_dict['builddependencies'].append(dep)
2091+
2092+
# ROCm-LLVM is handled separately: it is redistributable (kept as a runtime dep)
2093+
# and may be pulled in via a ROCm toolchain rather than as a direct dependency.
2094+
rocm_llvm_dep = find_rocm_llvm_dependency(ec)
2095+
if rocm_llvm_dep is not None:
2096+
add_gpu_property = 'add_property("arch","gpu")'
2097+
versionsuffix = rocm_llvm_dep['versionsuffix']
2098+
rocm_prefix = "-ROCm-"
2099+
2100+
if versionsuffix.startswith(rocm_prefix):
2101+
rocm_version = versionsuffix[len(rocm_prefix):]
2102+
else:
2103+
raise EasyBuildError(f"Invalid format for ROCm versionssuffix: {versionsuffix}")
2104+
pkg_versions['ROCm'] = rocm_version
2105+
20492106
if add_gpu_property:
20502107
ec.log.info("Injecting gpu as Lmod arch property and envvars for dependencies with their version")
20512108
modluafooter = 'modluafooter'

0 commit comments

Comments
 (0)