py-torch: set TORCH_CUDA_ARCH_LIST globally for dependents (#43962)

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 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
             "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 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
         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 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
         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 @@ class PyTorchaudio(PythonPackage):
 
         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 @@ class PyTorchvision(PythonPackage):
 
         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"]))