py-torch: add cusparselt variant (spack#1717)

* py-torch: add cusparselt variant

* py-nvidia-cusparselt: delete package

* [@spackbot] updating style on behalf of thomas-bouvier

* nvidia-cusparselt: disable cuda 13 for now

* [@spackbot] updating style on behalf of thomas-bouvier

* cusparselt: address reviews

---------

Co-authored-by: thomas-bouvier <[email protected]>

Author: thomas-bouvier

repos/spack_repo/builtin/packages/cusparselt/package.py

@@ -0,0 +1,66 @@
+# Copyright Spack Project Developers. See COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+import platform
+
+from spack.package import *
+
+
+class Cusparselt(Package):
+    """A high-performance CUDA library dedicated to general matrix-matrix operations
+    in which at least one operand is a structured sparse matrix with 50% sparsity ratio."""
+
+    homepage = "https://docs.nvidia.com/cuda/cusparselt/"
+
+    skip_version_audit = ["platform=darwin", "platform=windows"]
+
+    maintainers("thomas-bouvier")
+
+    system = platform.system().lower()
+    arch = platform.machine()
+    if "linux" in system and arch == "x86_64":
+        # version(
+        #    "0.8.1-cuda130",
+        #    sha256="82dd3e5ebc199a27011f58857a80cd825e77bba634ab2286ba3d4e13115db89a",
+        #    url="https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/libcusparse_lt-linux-x86_64-0.8.1.1_cuda13-archive.tar.xz",
+        # )
+        version(
+            "0.8.1-cuda120",
+            sha256="b34272e683e9f798435af05dc124657d1444cd0e13802c3d2f3152e31cd898a3",
+            url="https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/libcusparse_lt-linux-x86_64-0.8.1.1_cuda12-archive.tar.xz",
+        )
+    elif "linux" in system and arch == "aarch64":
+        # version(
+        #    "0.8.1-cuda130",
+        #    sha256="0fcf5808f66c71f755b4a73af2e955292e4334fec6a851eea1ac2e20878602b7",
+        #    url="https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-aarch64/libcusparse_lt-linux-aarch64-0.8.1.1_cuda13-archive.tar.xz",
+        # )
+        version(
+            "0.8.1-cuda120",
+            sha256="5426a897c73a9b98a83c4e132d15abc63dc4a00f7e38266e7b82c42cd58a01e1",
+            url="https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-aarch64/libcusparse_lt-linux-aarch64-0.8.1.1_cuda12-archive.tar.xz",
+        )
+
+    # cuda130_versions = ("@0.8.1-cuda130",)
+    cuda120_versions = ("@0.8.1-cuda120",)
+
+    # for v in cuda130_versions:
+    #    depends_on("cuda@13", when=v, type=("build", "run"))
+    for v in cuda120_versions:
+        depends_on("cuda@12", when=v, type=("build", "run"))
+
+    depends_on("c", type="build")
+    depends_on("cxx", type="build")
+
+    # Installation instructions
+    def install(self, spec, prefix):
+        # Create installation directories
+        mkdirp(prefix.lib)
+        mkdirp(prefix.include)
+
+        # Copy library files
+        install_tree("lib", prefix.lib)
+
+        # Copy header files
+        install_tree("include", prefix.include)

repos/spack_repo/builtin/packages/py_torch/package.py

@@ -110,6 +110,7 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     _desc = "Build the flash_attention kernel for scaled dot product attention"
     variant("flash_attention", default=True, description=_desc, when="@1.13:+cuda")
     variant("flash_attention", default=True, description=_desc, when="@1.13:+rocm")
+    variant("cusparselt", default=True, description="Use NVIDIA cuSPARSELt", when="@2.1: +cuda")
     # py-torch has strict dependencies on old protobuf/py-protobuf versions that
     # cause problems with other packages that require newer versions of protobuf
     # and py-protobuf --> provide an option to use the internal/vendored protobuf.
@@ -314,6 +315,7 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     depends_on("ucc", when="+ucc")
     depends_on("ucx", when="+ucc")
     depends_on("mkl", when="+mkldnn")
+    depends_on("cusparselt", when="+cusparselt")
 
     # Test dependencies
     with default_args(type="test"):
@@ -678,6 +680,7 @@ def enable_or_disable(variant, keyword="USE", var=None):
             env.set("CUDNN_INCLUDE_DIR", self.spec["cudnn"].prefix.include)
             env.set("CUDNN_LIBRARY", self.spec["cudnn"].libs[0])
 
+        enable_or_disable("cusparselt")
         enable_or_disable("fbgemm")
         enable_or_disable("kineto")
         enable_or_disable("magma")