diff --git a/easybuild/easyblocks/n/nccl.py b/easybuild/easyblocks/n/nccl.py index 48380466a94..b2ba096434f 100644 --- a/easybuild/easyblocks/n/nccl.py +++ b/easybuild/easyblocks/n/nccl.py @@ -31,7 +31,6 @@ import os from easybuild.easyblocks.generic.configuremake import ConfigureMake -from easybuild.tools.config import build_option from easybuild.tools.systemtools import get_shared_lib_ext from easybuild.tools.filetools import copy_file @@ -39,6 +38,11 @@ class EB_NCCL(ConfigureMake): """Support for building NCCL.""" + def __init__(self, *args, **kwargs): + """Init module env""" + super().__init__(*args, **kwargs) + self.module_load_environment.NCCL_HOME = [''] + def configure_step(self): """NCCL has no configure step""" pass @@ -47,17 +51,8 @@ def build_step(self): """Build NCCL""" # NCCL builds for all supported CUDA compute capabilities by default # If cuda_compute_capabilities is specified then we override this with the selected options - - # list of CUDA compute capabilities to use can be specifed in three ways (where 3 overrules 2 overrules 1): - # (1) in the easyconfig file, via the custom cuda_compute_capabilities; - # (2) via the EasyBuild environment variable EASYBUILD_CUDA_COMPUTE_CAPABILITIES; - # (3) in the EasyBuild configuration, via --cuda-compute-capabilities configuration option; - cuda_cc = build_option('cuda_compute_capabilities') or self.cfg['cuda_compute_capabilities'] - - nvcc_gencode = [] - for cc in cuda_cc: - add = cc.replace('.', '') - nvcc_gencode.append('-gencode=arch=compute_%s,code=sm_%s' % (add, add)) + cuda_cc = self.cfg.get_cuda_cc_template_value("cuda_cc_space_sep_no_period", required=False).split() + nvcc_gencode = [f'-gencode=arch=compute_{cc},code=sm_{cc}' for cc in cuda_cc] if nvcc_gencode: self.cfg.update('buildopts', 'NVCC_GENCODE="%s"' % ' '.join(nvcc_gencode))