Skip to content

Update jaxlib easyblock for jax 0.6.2 with CUDA-12.6.0 #3951

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
wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Update jaxlib easyblock for jax 0.6.2 with CUDA-12.6.0 #3951

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
51 changes: 39 additions & 12 deletions easybuild/easyblocks/j/jaxlib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,13 @@
@author: Denis Kristak (INUITS)
@author: Alexander Grund (TU Dresden)
@author: Alex Domingo (Vrije Universiteit Brussel)
@author: Pavel Tomanek (INUITS)
"""

import os
import tempfile

from easybuild.tools import LooseVersion

import easybuild.tools.environment as env
from easybuild.easyblocks.generic.pythonpackage import PythonPackage
from easybuild.framework.easyconfig import CUSTOM
Expand Down Expand Up @@ -106,31 +106,58 @@ def configure_step(self):
# Add optimization flags set by EasyBuild each as a separate option
bazel_options.extend(['--copt=%s' % i for i in os.environ['CXXFLAGS'].split(' ')])

# CUDA version
cuda_root = get_software_root('CUDA')
if cuda_root:
cudnn_root = get_software_root('cuDNN')
if not cudnn_root:
raise EasyBuildError('For CUDA-enabled builds cuDNN is also required')
nccl_root = get_software_root('NCCL')
cuda_version = '.'.join(get_software_version('CUDA').split('.')[:2]) # maj.minor
cudnn_version = '.'.join(get_software_version('cuDNN').split('.')[:3]) # maj.minor.patch
cuda_cc = self.cfg.get_cuda_cc_template_value('cuda_compute_capabilities')
options.extend([
'--enable_cuda',
'--cuda_path=' + cuda_root,
'--cuda_compute_capabilities=' + self.cfg.get_cuda_cc_template_value('cuda_compute_capabilities'),
'--cuda_compute_capabilities=' + cuda_cc,
'--cuda_version=' + cuda_version,
'--cudnn_path=' + cudnn_root,
'--cudnn_version=' + cudnn_version,
])

if LooseVersion(self.version) >= LooseVersion('0.1.70'):
nccl_root = get_software_root('NCCL')
if nccl_root:
options.append('--enable_nccl')
if LooseVersion(self.version) <= LooseVersion('0.4.33'):
options.extend([
'--enable_cuda',
'--cuda_path=' + cuda_root,
'--cudnn_path=' + cudnn_root,
])
if LooseVersion(self.version) >= LooseVersion('0.1.70'):
if nccl_root:
options.append('--enable_nccl')
else:
options.append('--noenable_nccl')
else: # from version 0.4.34 on
hermetic_cuda_cc = ','.join(f"sm_{cc.replace('.', '')}" for cc in cuda_cc.split(','))
bazel_options.extend([
f'--repo_env=HERMETIC_CUDA_VERSION={cuda_version}',
f'--repo_env=HERMETIC_CUDNN_VERSION={cudnn_version}',
f'--repo_env=LOCAL_CUDA_PATH={cuda_root}',
f'--repo_env=LOCAL_CUDNN_PATH={cudnn_root}',
f"--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES={hermetic_cuda_cc}",
*([f'--repo_env=LOCAL_NCCL_PATH={nccl_root}'] if nccl_root else []),
])
# set Clang flags - CUDA version needs Clang to be built
clang_root = get_software_root('Clang')
if clang_root:
options.extend([
'--use_clang=true',
f'--clang_path={os.path.join(clang_root, "bin", "clang++")}',
])
bazel_options.extend([
'--@local_config_cuda//:cuda_compiler=clang',
'--@local_config_cuda//cuda:include_cuda_libs=true',
])
else:
options.append('--noenable_nccl')
raise EasyBuildError('For CUDA-enabled builds Clang is also required')

config_env_vars['GCC_HOST_COMPILER_PATH'] = which(os.getenv('CC'))
elif LooseVersion(self.version) <= LooseVersion('0.6.0'):
elif LooseVersion(self.version) <= LooseVersion('0.4.33'):
options.append('--noenable_cuda')

if self.cfg['use_mkl_dnn']:
Expand Down