45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
# Copyright 2013-2022 Lawrence Livermore National Security, LLC and other
|
|
# Spack Project Developers. See the top-level COPYRIGHT file for details.
|
|
#
|
|
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
|
|
|
|
|
|
from spack.package import *
|
|
|
|
|
|
class PyTorchNvidiaApex(PythonPackage, CudaPackage):
|
|
"""A PyTorch Extension: Tools for easy mixed precision and
|
|
distributed training in Pytorch"""
|
|
|
|
homepage = "https://github.com/nvidia/apex/"
|
|
git = "https://github.com/nvidia/apex/"
|
|
|
|
version("master", branch="master")
|
|
version("2020-10-19", commit="8a1ed9e8d35dfad26fb973996319965e4224dcdd")
|
|
|
|
depends_on("python@3:", type=("build", "run"))
|
|
depends_on("py-setuptools", type="build")
|
|
depends_on("py-torch@0.4:", type=("build", "run"))
|
|
depends_on("cuda@9:", when="+cuda")
|
|
depends_on("py-pybind11", type=("build", "link", "run"))
|
|
|
|
variant("cuda", default=True, description="Build with CUDA")
|
|
|
|
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)
|
|
|
|
def install_options(self, spec, prefix):
|
|
args = []
|
|
if spec.satisfies("^py-torch@1.0:"):
|
|
args.append("--cpp_ext")
|
|
if "+cuda" in spec:
|
|
args.append("--cuda_ext")
|
|
return args
|