nccl-tests: respect cuda_arch variant during build (#21995)

2021-02-26 18:38:20 -05:00 · 2021-02-26 18:38:20 -05:00 · f2f58b70b1
commit f2f58b70b1
parent f957f1936c
1 changed files with 5 additions and 0 deletions
--- a/var/spack/repos/builtin/packages/nccl-tests/package.py
+++ b/var/spack/repos/builtin/packages/nccl-tests/package.py
@ -31,6 +31,11 @@ def build_targets(self):
            targets.append('MPI=1')
        return targets

+    def setup_build_environment(self, env):
+        cuda_arch = self.spec.variants["cuda_arch"].value
+        cuda_gencode = " ".join(self.cuda_flags(cuda_arch))
+        env.set("NVCC_GENCODE", cuda_gencode)
+
    def install(self, spec, prefix):
        mkdirp(prefix.bin)
        install_tree('./build', prefix.bin)