diff --git a/pkgs/development/cuda-modules/overrides/cublasmp/libcublasmp.nix b/pkgs/development/cuda-modules/overrides/cublasmp/libcublasmp.nix index 5099a316fa8630f..4d8e0f986db7afc 100644 --- a/pkgs/development/cuda-modules/overrides/cublasmp/libcublasmp.nix +++ b/pkgs/development/cuda-modules/overrides/cublasmp/libcublasmp.nix @@ -1,5 +1,13 @@ -{ libcublas, libcal }: +args@{ + lib, + libcal ? null, + libcublas ? null, + utils, +}: prevAttrs: { + brokenConditions = prevAttrs.brokenConditions // (utils.mkMissingPackagesBrokenConditions args); + # TODO: Looks like the minimum supported capability is 7.0 as of the latest: + # https://docs.nvidia.com/cuda/cublasmp/getting_started/index.html buildInputs = prevAttrs.buildInputs or [ ] ++ [ libcal.lib libcublas.lib diff --git a/pkgs/development/cuda-modules/overrides/cuquantum/cuquantum.nix b/pkgs/development/cuda-modules/overrides/cuquantum/cuquantum.nix index b0bb4e11b8609f9..245adca15a7cb84 100644 --- a/pkgs/development/cuda-modules/overrides/cuquantum/cuquantum.nix +++ b/pkgs/development/cuda-modules/overrides/cuquantum/cuquantum.nix @@ -1,12 +1,16 @@ -{ - libcutensor, - libcublas, - libcusolver, +args@{ + cudaOlder, + lib, + libcublas ? null, + libcusolver ? null, + libcutensor ? null, + utils, }: prevAttrs: { + brokenConditions = prevAttrs.brokenConditions // (utils.mkMissingPackagesBrokenConditions args); buildInputs = prevAttrs.buildInputs or [ ] ++ [ - libcutensor.lib libcublas.lib libcusolver.lib + libcutensor.lib ]; } diff --git a/pkgs/development/cuda-modules/overrides/cusolvermp/libcusolvermp.nix b/pkgs/development/cuda-modules/overrides/cusolvermp/libcusolvermp.nix index 0317193a14eaf50..f35cc6d482cadf5 100644 --- a/pkgs/development/cuda-modules/overrides/cusolvermp/libcusolvermp.nix +++ b/pkgs/development/cuda-modules/overrides/cusolvermp/libcusolvermp.nix @@ -1,10 +1,13 @@ -{ - libcublas, - libcal, - libcusolver, - cuda_cudart, +args@{ + cuda_cudart ? null, + lib, + libcal ? null, + libcublas ? null, + libcusolver ? null, + utils, }: prevAttrs: { + brokenConditions = prevAttrs.brokenConditions // (utils.mkMissingPackagesBrokenConditions args); buildInputs = prevAttrs.buildInputs or [ ] ++ [ libcal.lib libcublas.lib diff --git a/pkgs/development/cuda-modules/package-sets.nix b/pkgs/development/cuda-modules/package-sets.nix index 0ec610caf613d7b..598cc31e14980bb 100644 --- a/pkgs/development/cuda-modules/package-sets.nix +++ b/pkgs/development/cuda-modules/package-sets.nix @@ -13,7 +13,7 @@ let stdenv ; inherit (lib.attrsets) - attrNames + attrByPath dontRecurseIntoAttrs filterAttrs mapAttrs @@ -134,31 +134,51 @@ let ... }: let - # There is a variant for the desired CUDA version. - isDesiredCudaVariant = cudaVariant == (utils.mkCudaVariant cudaMajorMinorPatchVersion); # One of the subdirectories of the lib directory contains a supported version for our version of CUDA. + # This is typically found with older versions of redistributables which don't use separate tarballs for each + # supported CUDA version. hasSupportedCudaVersionInLib = (utils.getLibPath cudaMajorMinorPatchVersion packageInfo.feature.cudaVersionsInLib) != null; + # There is a variant for the desired CUDA version. + isDesiredCudaVariant = cudaVariant == (utils.mkCudaVariant cudaMajorMinorPatchVersion); in - if redistName == "cuda" then + attrByPath [ redistName ] (isDesiredCudaVariant || hasSupportedCudaVersionInLib) { + # CUBLASMP: Looks like it requires at least 11.8: + # https://docs.nvidia.com/cuda/cublasmp/getting_started/index.html + cublasmp = versionAtLeast cudaMajorMinorPatchVersion "11.8"; + # CUDA: None of the CUDA redistributables have CUDA variants, but we only need to check that the release # version matches the CUDA version we want. - version == cudaMajorMinorPatchVersion - else if redistName == "cudnn" then + cuda = version == cudaMajorMinorPatchVersion; + # CUDNN: Since cuDNN 8.5, it is possible to use the dynamic library for a CUDA release with any CUDA version # in that major release series. For example, the cuDNN 8.5 dynamic library for CUDA 11.0 can be used with # any CUDA 11.x release. (This functionality is not present for the CUDA 10.2 releases.) # As such, it is enough that the cuda variant matches to accept the package. - isDesiredCudaVariant - else if redistName == "cutensor" then + cudnn = isDesiredCudaVariant; + + # CUQUANTUM: Only available for CUDA 11.5 and later. + cuquantum = + if cudaVariant == "None" then + # This handles the case of pre-23.03 releases, which don't provide CUDA versions in lib + versionAtLeast cudaMajorMinorPatchVersion "11.5" + # And this handles the case of pre-23.06 releases, which do + || hasSupportedCudaVersionInLib + else + isDesiredCudaVariant; + # CUTENSOR: Instead of providing CUDA variants, cuTensor provides multiple versions of the library nested # in the lib directory. So long as one of the versions in cudaVersionsInLib is a prefix of the current CUDA # version, we accept the package. We should have a more stringent version check, but no one has written # a sidecar file mapping releases to supported CUDA versions. - hasSupportedCudaVersionInLib - else - # TODO: - (isDesiredCudaVariant || hasSupportedCudaVersionInLib); + cutensor = hasSupportedCudaVersionInLib; + + # TODO: Add constraints for TensorRT + + # TODO: These constraints are duplicated in the overrides files because the overrides package set is created + # with callPackage and evaluated strictly. As such, the constraints should exist outside of either of these + # places -- perhaps in config.data? + }; # Function to determine if a package should be accepted into the package set. packagePlatformMatches = diff --git a/pkgs/development/cuda-modules/utils.nix b/pkgs/development/cuda-modules/utils.nix index a6d7f2ca2bd4cda..c88c980143fd53f 100644 --- a/pkgs/development/cuda-modules/utils.nix +++ b/pkgs/development/cuda-modules/utils.nix @@ -1,6 +1,11 @@ { config, lib, ... }: let - inherit (lib.attrsets) attrByPath mapAttrs; + inherit (lib.attrsets) + attrByPath + filterAttrs + mapAttrs + mapAttrs' + ; inherit (lib.filesystem) packagesFromDirectoryRecursive; inherit (lib.lists) findFirst @@ -18,6 +23,7 @@ let ; inherit (lib.types) attrsOf + bool enum functionTo nonEmptyListOf @@ -159,6 +165,23 @@ in (builtins.removeAttrs args [ "leaf" ]) // { packageInfo = leaf; } ); }; + mkMissingPackagesBrokenConditions = { + description = '' + Utility function to generate a set of brokenConditions for missing packages. + ''; + type = functionTo (attrsOf bool); + default = flip pipe [ + # Take the attributes that are null. + (filterAttrs (_: value: value == null)) + # Map them to a set of brokenConditions. + (mapAttrs' ( + name: value: { + name = "Required package ${name} is missing"; + value = true; + } + )) + ]; + }; mkTrimmedFlattenedIndex = { description = '' Helper function to reduce the number of attribute sets we need to traverse to create all the package sets.