From 14561fafffbcf311775f8a4844a8b0faffaed40c Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Fri, 3 May 2024 09:11:59 +0200 Subject: [PATCH] py-torch: set TORCH_CUDA_ARCH_LIST globally for dependents (#43962) --- .../builtin/packages/py-torch-cluster/package.py | 5 ----- .../packages/py-torch-geometric/package.py | 5 ----- .../packages/py-torch-nvidia-apex/package.py | 6 ------ .../builtin/packages/py-torch-scatter/package.py | 5 ----- .../builtin/packages/py-torch-sparse/package.py | 5 ----- .../packages/py-torch-spline-conv/package.py | 5 ----- .../repos/builtin/packages/py-torch/package.py | 16 ++++++++++++---- .../builtin/packages/py-torchaudio/package.py | 5 ----- .../builtin/packages/py-torchvision/package.py | 5 ----- 9 files changed, 12 insertions(+), 45 deletions(-) diff --git a/var/spack/repos/builtin/packages/py-torch-cluster/package.py b/var/spack/repos/builtin/packages/py-torch-cluster/package.py index feb225ff9a..685019395f 100644 --- a/var/spack/repos/builtin/packages/py-torch-cluster/package.py +++ b/var/spack/repos/builtin/packages/py-torch-cluster/package.py @@ -35,11 +35,6 @@ class PyTorchCluster(PythonPackage): def setup_build_environment(self, env): if "+cuda" in self.spec: - cuda_arches = list(self.spec["py-torch"].variants["cuda_arch"].value) - for i, x in enumerate(cuda_arches): - cuda_arches[i] = "{0}.{1}".format(x[0:-1], x[-1]) - env.set("TORCH_CUDA_ARCH_LIST", str.join(" ", cuda_arches)) - env.set("FORCE_CUDA", "1") env.set("CUDA_HOME", self.spec["cuda"].prefix) else: diff --git a/var/spack/repos/builtin/packages/py-torch-geometric/package.py b/var/spack/repos/builtin/packages/py-torch-geometric/package.py index 89b5ffe181..e9820c2bdc 100644 --- a/var/spack/repos/builtin/packages/py-torch-geometric/package.py +++ b/var/spack/repos/builtin/packages/py-torch-geometric/package.py @@ -72,11 +72,6 @@ class PyTorchGeometric(PythonPackage): def setup_build_environment(self, env): if "+cuda" in self.spec: - cuda_arches = list(self.spec["py-torch"].variants["cuda_arch"].value) - for i, x in enumerate(cuda_arches): - cuda_arches[i] = "{0}.{1}".format(x[0:-1], x[-1]) - env.set("TORCH_CUDA_ARCH_LIST", str.join(" ", cuda_arches)) - env.set("FORCE_CUDA", "1") env.set("CUDA_HOME", self.spec["cuda"].prefix) else: diff --git a/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py b/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py index 00a357603f..08c0ecad6d 100644 --- a/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py +++ b/var/spack/repos/builtin/packages/py-torch-nvidia-apex/package.py @@ -35,12 +35,6 @@ class PyTorchNvidiaApex(PythonPackage, CudaPackage): def setup_build_environment(self, env): if "+cuda" in self.spec: env.set("CUDA_HOME", self.spec["cuda"].prefix) - if self.spec.variants["cuda_arch"].value[0] != "none": - torch_cuda_arch = ";".join( - "{0:.1f}".format(float(i) / 10.0) - for i in self.spec.variants["cuda_arch"].value - ) - env.set("TORCH_CUDA_ARCH_LIST", torch_cuda_arch) else: env.unset("CUDA_HOME") diff --git a/var/spack/repos/builtin/packages/py-torch-scatter/package.py b/var/spack/repos/builtin/packages/py-torch-scatter/package.py index 52804ef4b6..e1aea29192 100644 --- a/var/spack/repos/builtin/packages/py-torch-scatter/package.py +++ b/var/spack/repos/builtin/packages/py-torch-scatter/package.py @@ -29,11 +29,6 @@ class PyTorchScatter(PythonPackage): def setup_build_environment(self, env): if "+cuda" in self.spec: - cuda_arches = list(self.spec["py-torch"].variants["cuda_arch"].value) - for i, x in enumerate(cuda_arches): - cuda_arches[i] = "{0}.{1}".format(x[0:-1], x[-1]) - env.set("TORCH_CUDA_ARCH_LIST", str.join(" ", cuda_arches)) - env.set("FORCE_CUDA", "1") env.set("CUDA_HOME", self.spec["cuda"].prefix) else: diff --git a/var/spack/repos/builtin/packages/py-torch-sparse/package.py b/var/spack/repos/builtin/packages/py-torch-sparse/package.py index 8a1ed1a57d..8df9e1bd62 100644 --- a/var/spack/repos/builtin/packages/py-torch-sparse/package.py +++ b/var/spack/repos/builtin/packages/py-torch-sparse/package.py @@ -31,11 +31,6 @@ class PyTorchSparse(PythonPackage): def setup_build_environment(self, env): if "+cuda" in self.spec: - cuda_arches = list(self.spec["py-torch"].variants["cuda_arch"].value) - for i, x in enumerate(cuda_arches): - cuda_arches[i] = "{0}.{1}".format(x[0:-1], x[-1]) - env.set("TORCH_CUDA_ARCH_LIST", str.join(" ", cuda_arches)) - env.set("FORCE_CUDA", "1") env.set("CUDA_HOME", self.spec["cuda"].prefix) else: diff --git a/var/spack/repos/builtin/packages/py-torch-spline-conv/package.py b/var/spack/repos/builtin/packages/py-torch-spline-conv/package.py index e1bc630e54..b861734b49 100644 --- a/var/spack/repos/builtin/packages/py-torch-spline-conv/package.py +++ b/var/spack/repos/builtin/packages/py-torch-spline-conv/package.py @@ -27,11 +27,6 @@ class PyTorchSplineConv(PythonPackage): def setup_build_environment(self, env): if "+cuda" in self.spec: - cuda_arches = list(self.spec["py-torch"].variants["cuda_arch"].value) - for i, x in enumerate(cuda_arches): - cuda_arches[i] = "{0}.{1}".format(x[0:-1], x[-1]) - env.set("TORCH_CUDA_ARCH_LIST", str.join(" ", cuda_arches)) - env.set("FORCE_CUDA", "1") env.set("CUDA_HOME", self.spec["cuda"].prefix) else: diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py index 1dff8b12e0..5de618494d 100644 --- a/var/spack/repos/builtin/packages/py-torch/package.py +++ b/var/spack/repos/builtin/packages/py-torch/package.py @@ -473,6 +473,13 @@ def patch(self): "caffe2/CMakeLists.txt", ) + def torch_cuda_arch_list(self, env): + if "+cuda" in self.spec: + torch_cuda_arch = ";".join( + "{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value + ) + env.set("TORCH_CUDA_ARCH_LIST", torch_cuda_arch) + def setup_build_environment(self, env): """Set environment variables used to control the build. @@ -515,10 +522,8 @@ def enable_or_disable(variant, keyword="USE", var=None): if "+cuda" in self.spec: env.set("CUDA_HOME", self.spec["cuda"].prefix) # Linux/macOS env.set("CUDA_PATH", self.spec["cuda"].prefix) # Windows - torch_cuda_arch = ";".join( - "{0:.1f}".format(float(i) / 10.0) for i in self.spec.variants["cuda_arch"].value - ) - env.set("TORCH_CUDA_ARCH_LIST", torch_cuda_arch) + self.torch_cuda_arch_list(env) + if self.spec.satisfies("%clang"): for flag in self.spec.compiler_flags["cxxflags"]: if "gcc-toolchain" in flag: @@ -667,6 +672,9 @@ def enable_or_disable(variant, keyword="USE", var=None): if self.spec.satisfies("%apple-clang@15:"): env.append_flags("LDFLAGS", "-Wl,-ld_classic") + def setup_run_environment(self, env): + self.torch_cuda_arch_list(env) + @run_before("install") def build_amd(self): if "+rocm" in self.spec: diff --git a/var/spack/repos/builtin/packages/py-torchaudio/package.py b/var/spack/repos/builtin/packages/py-torchaudio/package.py index e877afc10c..d1fa9f4f78 100644 --- a/var/spack/repos/builtin/packages/py-torchaudio/package.py +++ b/var/spack/repos/builtin/packages/py-torchaudio/package.py @@ -102,11 +102,6 @@ def setup_build_environment(self, env): if "+cuda" in self.spec["py-torch"]: env.set("USE_CUDA", 1) - torch_cuda_arch_list = ";".join( - "{0:.1f}".format(float(i) / 10.0) - for i in self.spec["py-torch"].variants["cuda_arch"].value - ) - env.set("TORCH_CUDA_ARCH_LIST", torch_cuda_arch_list) else: env.set("USE_CUDA", 0) diff --git a/var/spack/repos/builtin/packages/py-torchvision/package.py b/var/spack/repos/builtin/packages/py-torchvision/package.py index 08e6182300..96994a5b7a 100644 --- a/var/spack/repos/builtin/packages/py-torchvision/package.py +++ b/var/spack/repos/builtin/packages/py-torchvision/package.py @@ -150,11 +150,6 @@ def setup_build_environment(self, env): if "^cuda" in self.spec: env.set("CUDA_HOME", self.spec["cuda"].prefix) - torch_cuda_arch_list = ";".join( - "{0:.1f}".format(float(i) / 10.0) - for i in self.spec["py-torch"].variants["cuda_arch"].value - ) - env.set("TORCH_CUDA_ARCH_LIST", torch_cuda_arch_list) for gpu in ["cuda", "mps"]: env.set(f"FORCE_{gpu.upper()}", int(f"+{gpu}" in self.spec["py-torch"]))