From 06f4f4391542a46fc7404c1e85a574b464a3a941 Mon Sep 17 00:00:00 2001 From: Brian Van Essen Date: Thu, 30 Jun 2022 22:37:47 -0700 Subject: New package: NVIDIA cuTensor (#31397) --- .../repos/builtin/packages/cutensor/package.py | 57 ++++++++++++++++++++++ var/spack/repos/builtin/packages/lbann/package.py | 1 + 2 files changed, 58 insertions(+) create mode 100644 var/spack/repos/builtin/packages/cutensor/package.py (limited to 'var') diff --git a/var/spack/repos/builtin/packages/cutensor/package.py b/var/spack/repos/builtin/packages/cutensor/package.py new file mode 100644 index 0000000000..14734bd922 --- /dev/null +++ b/var/spack/repos/builtin/packages/cutensor/package.py @@ -0,0 +1,57 @@ +# 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) + +import platform + +from spack.package import * + +_versions = { + # cuTensor 1.5.0 + '1.5.0.3': { + 'Linux-x86_64': '4fdebe94f0ba3933a422cff3dd05a0ef7a18552ca274dd12564056993f55471d', + 'Linux-ppc64le': 'ad736acc94e88673b04a3156d7d3a408937cac32d083acdfbd8435582cbe15db', + 'Linux-aarch64': '5b9ac479b1dadaf40464ff3076e45f2ec92581c07df1258a155b5bcd142f6090'}, + +} + + +class Cutensor(Package): + """NVIDIA cuTENSOR Library is a GPU-accelerated tensor linear algebra + library providing tensor contraction, reduction and elementwise + operations.""" + + homepage = "https://developer.nvidia.com/cutensor" + + maintainers = ['bvanessen'] + url = "cutensor" + + for ver, packages in _versions.items(): + key = "{0}-{1}".format(platform.system(), platform.machine()) + pkg = packages.get(key) + cutensor_ver = ver + + cuda_ver = '10.0' + if platform.machine() == 'aarch64': + cuda_ver = '11.0' + + if pkg: + version(cutensor_ver, sha256=pkg) + # Add constraints matching CUDA version to cuTensor version + cuda_req = 'cuda@{0}:'.format(cuda_ver) + cutensor_ver_req = '@{0}'.format(cutensor_ver) + depends_on(cuda_req, when=cutensor_ver_req) + + def url_for_version(self, version): + # Get the system and machine arch for building the file path + sys = "{0}-{1}".format(platform.system(), platform.machine()) + # Munge it to match Nvidia's naming scheme + sys_key = sys.lower() + sys_key = sys_key.replace('aarch64', 'sbsa') + + url = 'https://developer.download.nvidia.com/compute/cutensor/redist/libcutensor/{0}/libcutensor-{0}-{1}-archive.tar.xz' + return url.format(sys_key, version) + + def install(self, spec, prefix): + install_tree('.', prefix) diff --git a/var/spack/repos/builtin/packages/lbann/package.py b/var/spack/repos/builtin/packages/lbann/package.py index 6c80069802..3500d3eb62 100644 --- a/var/spack/repos/builtin/packages/lbann/package.py +++ b/var/spack/repos/builtin/packages/lbann/package.py @@ -172,6 +172,7 @@ class Lbann(CMakePackage, CudaPackage, ROCmPackage): depends_on('cudnn', when='@0.90:0.100 +cuda') depends_on('cudnn@8.0.2:', when='@:0.90,0.101: +cuda') depends_on('cub', when='@0.94:0.98.2 +cuda ^cuda@:10') + depends_on('cutensor', when='@:0.90,0.102: +cuda') depends_on('hipcub', when='+rocm') depends_on('mpi') depends_on('hwloc@1.11:', when='@:0.90,0.102: +hwloc') -- cgit v1.2.3-60-g2f50