Skip to content

Add gpu tag for ROCm-LLVM and refactor method #228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zerefwayne wants to merge 1 commit into
base: main
Choose a base branch
from
+27 −18
Open

Add gpu tag for ROCm-LLVM and refactor method #228

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 27 additions & 18 deletions eb_hooks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1944,34 +1944,43 @@ def inject_gpu_property(ec):
# - drop dependency to build dependency
# - add 'gpu' Lmod property
# - add envvar with package version
pkg_names = ( "CUDA", "cuDNN" )
pkg_names = ( "CUDA", "cuDNN", "ROCm-LLVM" )
# skip dropping step for ROCm-LLVM as it is redistributable
drop_to_builddep = ( "CUDA", "cuDNN" )
pkg_versions = { }
add_gpu_property = ''

# If none of the gpu packages are in the easyconfig, do not process further
dep_names = {dep[0] for dep in ec_dict['dependencies']}
Copy link
Copy Markdown
Member

@ocaisa ocaisa May 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you tested this? I wonder if ROCm-LLVM is the toolchain if it actually appears in the dependencies?

Copy link
Copy Markdown
Contributor

@casparvl casparvl May 12, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I doubt it. For one of the other hooks (can't remember which) where we searched some dep list with CUDA / cuDNN in it, I added logic to also check the toolchain instead - and then just trigger it if it's rocm-compilers, rompi,... etc. This should probably do something similar. @zerefwayne search for rocm-compilers in eb_hooks.py and you'll probably find my logic in the other hook.

Copy link
Copy Markdown
Member

@ocaisa ocaisa May 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be better if we can identify ROCm-LLVM explicitly so that we can store the version in an environment variable. That will be useful later in Lmod when we are checking if the version is supported by the driver.

if not any(pkg in dep_names for pkg in pkg_names):
return ec

# Check if pkg_name is in the dependencies, if so drop dependency to build
# dependency and set variable for later adding the 'gpu' Lmod property
# to '.remove' dependencies from ec_dict['dependencies'] we make a copy,
# iterate over the copy and can then savely use '.remove' on the original
for pkg_name in pkg_names:
# Check if pkg_name is in the dependencies, if so drop dependency to build
# dependency and set variable for later adding the 'gpu' Lmod property
# to '.remove' dependencies from ec_dict['dependencies'] we make a copy,
# iterate over the copy and can then savely use '.remove' on the original
# ec_dict['dependencies'].
deps = ec_dict['dependencies'][:]
if (pkg_name in [dep[0] for dep in deps]):
for dep in ec_dict['dependencies'][:]:
if dep[0] != pkg_name:
continue

add_gpu_property = 'add_property("arch","gpu")'
for dep in deps:
if pkg_name == dep[0]:
# make pkg_name a build dependency only (rpathing saves us from link errors)
ec.log.info("Dropping dependency on %s to build dependency" % pkg_name)
ec_dict['dependencies'].remove(dep)
if dep not in ec_dict['builddependencies']:
ec_dict['builddependencies'].append(dep)
# take note of version for creating the modluafooter
pkg_versions[pkg_name] = dep[1]
pkg_versions[pkg_name] = dep[1]

if pkg_name in drop_to_builddep:
ec.log.info("Dropping dependency on %s to build dependency" % pkg_name)
ec_dict['dependencies'].remove(dep)
# Avoid adding a duplicate
if dep not in ec_dict['builddependencies']:
ec_dict['builddependencies'].append(dep)

if add_gpu_property:
ec.log.info("Injecting gpu as Lmod arch property and envvars for dependencies with their version")
modluafooter = 'modluafooter'
extra_mod_footer_lines = [add_gpu_property]
for pkg_name, version in pkg_versions.items():
envvar = "EESSI%sVERSION" % pkg_name.upper()
# follow the convention for replacing - with MIN
envvar = "EESSI%sVERSION" % pkg_name.upper().replace('-', 'MIN')
extra_mod_footer_lines.append('setenv("%s","%s")' % (envvar, version))
# take into account that modluafooter may already be set
if modluafooter in ec_dict:
Expand Down
Loading