diff --git a/pkgs/development/cuda-modules/package-sets.nix b/pkgs/development/cuda-modules/package-sets.nix index 11cf95de3d9e2c3..d82368bd60ed398 100644 --- a/pkgs/development/cuda-modules/package-sets.nix +++ b/pkgs/development/cuda-modules/package-sets.nix @@ -6,22 +6,10 @@ }: let inherit (config) data utils; - inherit (pkgs) - fetchurl - newScope - srcOnly - stdenv - ; - inherit (lib.attrsets) - attrByPath - dontRecurseIntoAttrs - filterAttrs - mapAttrs - optionalAttrs - ; + inherit (pkgs) newScope stdenv; + inherit (lib.attrsets) dontRecurseIntoAttrs mapAttrs optionalAttrs; inherit (lib.customisation) makeScope; inherit (lib.filesystem) packagesFromDirectoryRecursive; - inherit (lib.licenses) nvidiaCudaRedist; inherit (lib.lists) foldl' map @@ -29,10 +17,9 @@ let optionals unique ; - inherit (lib.meta) addMetaAttrs; inherit (lib.options) mkOption; - inherit (lib.strings) removeSuffix versionAtLeast versionOlder; - inherit (lib.trivial) const flip pipe; + inherit (lib.strings) versionAtLeast versionOlder; + inherit (lib.trivial) const pipe; inherit (lib.types) attrsOf raw; inherit (lib.versions) major majorMinor; @@ -41,151 +28,12 @@ let # - Overrides for cutensor, etc. # - Rename the platform type to redistPlatform to distinguish it from the Nix platform. - redistArch = utils.getRedistArch stdenv.hostPlatform.system; - isSupportedRedistArch = redistArch != "unsupported"; - - # Function to override a package. - overridePackage = - final: - let - # Maps redist name to package name to override. - overrides = utils.getOverrides final; - in - { packageName, redistName, ... }: - package: package.overrideAttrs (overrides.${redistName}.${packageName} or { }); - - getNixPlatforms = - platform: - if platform == "source" then - # All platforms are supported for source packages. - [ - "aarch64-linux" # Both SBSA (ARM servers) and Jetson - "x86_64-linux" # x86_64 - "ppc64le-linux" # POWER - ] - else - [ (utils.getNixPlatform platform) ]; - - # Function to build a redistributable package. - buildRedistPackage = - final: cudaMajorMinorPatchVersion: - meta@{ - packageName, - redistName, - releaseInfo, - packageInfo, - version, - ... - }: - let - package = pipe ./manifest-builder.nix [ - # Build the package. - (flip final.callPackage { - inherit (meta) packageInfo packageName releaseInfo; - libPath = utils.getLibPath cudaMajorMinorPatchVersion packageInfo.feature.cudaVersionsInLib; - # The source is given by the tarball, which we unpack and use as a FOD. - src = - let - tarball = fetchurl { - inherit (packageInfo) sha256; - url = - if redistName == "tensorrt" then - utils.mkTensorRTURL version packageInfo.relativePath - else - utils.mkRedistURL redistName (utils.mkRelativePath meta); - }; - unpacked = srcOnly { - __structuredAttrs = true; - strictDeps = true; - name = tarball.name + "-unpacked"; - src = tarball; - outputHashMode = "recursive"; - outputHash = packageInfo.narHash; - }; - in - unpacked; - }) - # Update the package license - (addMetaAttrs { - license = nvidiaCudaRedist // { - url = - let - licensePath = - if releaseInfo.licensePath != null then releaseInfo.licensePath else "${packageName}/LICENSE.txt"; - in - "https://developer.download.nvidia.com/compute/${redistName}/redist/${licensePath}"; - }; - }) - # Apply package-specific overrides, if they exist. - (overridePackage final meta) - ]; - in - package; - - # Function to determine if a package should be accepted into the package set. - packageCudaVariantMatches = - cudaMajorMinorPatchVersion: - { - cudaVariant, - packageInfo, - platform, - redistName, - version, - ... - }: - let - # 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 - 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. - 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. - 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. - 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 = - platform: platform == "source" || (isSupportedRedistArch && platform == redistArch); - # Function to update the package set with the new package. updatePackages = + let + redistArch = utils.getRedistArch stdenv.hostPlatform.system; + isSupportedRedistArch = redistArch != "unsupported"; + in { packageInfo, packageName, @@ -194,7 +42,6 @@ let releaseInfo, ... }: - packages: package: let # Non-CUDA redist packages, like cuDNN, are multiplexed so we add a suffix to the attribute name. # CUDA redist packages are always the same as the package name. @@ -204,25 +51,39 @@ let else packageName; - # We want to make sure that packages exist even on platforms they cannot be built on or for, to avoid missing - # attribute errors. - # - # For each package we see, we do the following: - # - # - If a package of the same name does not exist in the package set: - # - If the package we see is for the current platform: - # - Add the package to the package set. - # - If the package we see is not for the current platform: - # - Add the package to the package set, but with a null `src` attribute and a `meta.platforms` attribute - # equal to the platform the package is for. - # - If a package of the same name exists in the package set: - # - If the package we see is for the current platform: - # - Override the `src` attribute and append the current platform to the `meta.platforms` attribute. - # - If the package we see is not for the current platform: - # - Append the current platform to the `meta.platforms` attribute. + # Package platform matches the current platform or the package is a source package. + packagePlatformMatches = platform == "source" || (isSupportedRedistArch && platform == redistArch); + in + packages: package: + let packageForName = name: let + # # We want to make sure that packages exist even on platforms they cannot be built on or for, to avoid missing + # # attribute errors. + # # + # # For each package we see, we do the following: + # # + # # - If a package of the same name exists in the package set: (packageExists) + # packageExists = packages ? ${name}; + # # - If the package we see is for the current platform: (packagePlatformMatches) + # # - Override the `src` attribute and append the current platform to the `meta.platforms` attribute. + # setSrcAndAppendPlatformToMetaPlatforms = packageExists && packagePlatformMatches; + # # - If the package we see is not for the current platform: (!packagePlatformMatches) + # # - Append the current platform to the `meta.platforms` attribute. (appendPlatformToMetaPlatforms) + # appendPlatformToMetaPlatforms = packageExists && !packagePlatformMatches; + # # - If a package of the same name does not exist in the package set: (!packageExists) + # # - If the package we see is for the current platform: (packagePlatformMatches) + # # - Add the package to the package set. (addPackageToPackageSet) + # addPackageToPackageSet = !packageExists && packagePlatformMatches; + # # - If the package we see is not for the current platform: (!packagePlatformMatches) + # # - Add the package to the package set, but with a null `src` attribute and a `meta.platforms` attribute + # # equal to the platform the package is for. (setSrcToNullAndMetaPlatformsToPlatform) + # setSrcToNullAndMetaPlatformsToPlatform = noPackageExists && !packagePlatformMatches; + + # # A package exists in the package set, but it's older than the package we're processing. + # existingPackageIsOlder = packages ? ${name} && versionOlder packages.${name}.version package.version; + replacePackage = # There's no entry for the package in the package set. !(packages ? ${name}) @@ -230,8 +91,6 @@ let || (versionOlder packages.${name}.version package.version); package' = if replacePackage then package else packages.${name}; - - platformMatched = packagePlatformMatches platform; in pipe package' [ # Update the package. @@ -240,16 +99,16 @@ let # Only do the override if the platform matched or we're replacing the package. # Otherwise leave it alone. pkg.override ( - optionalAttrs (platformMatched || replacePackage) { + optionalAttrs (packagePlatformMatches || replacePackage) { src = - if platformMatched then + if packagePlatformMatches then package.src else if replacePackage then null else builtins.throw "This should never happen."; packageInfo = - if platformMatched then + if packagePlatformMatches then packageInfo else if replacePackage then { feature.outputs = [ "out" ]; } @@ -265,7 +124,7 @@ let meta.platforms = let existingPlatforms = optionals (!replacePackage) (prevAttrs.meta.platforms or [ ]); - newPlatforms = getNixPlatforms platform; + newPlatforms = utils.getNixPlatforms platform; in # TODO: This has at least O(n^2) complexity due to the `unique` call. # And we do this for every package we process... @@ -304,15 +163,17 @@ let ... }: let - package = buildRedistPackage final cudaMajorMinorPatchVersion meta; + package = utils.buildRedistPackage final meta; in - if packageCudaVariantMatches cudaMajorMinorPatchVersion meta then + if utils.packageSupportsCudaVersion cudaMajorMinorPatchVersion meta then updatePackages meta packages package else packages; in { lib = dontRecurseIntoAttrs lib; + # TODO: pkgs.cudaPackages_11_8.pkgs.cudaPackages.cudaVersion should be 11.8, not pkgs.cudaPackages.cudaVersion. + # pkgs = dontRecurseIntoAttrs pkgs // { inherit (final) cudaPackages; }; pkgs = dontRecurseIntoAttrs pkgs; data = dontRecurseIntoAttrs config.data; utils = dontRecurseIntoAttrs config.utils; diff --git a/pkgs/development/cuda-modules/utils.nix b/pkgs/development/cuda-modules/utils.nix index c88c980143fd53f..72b46048efa2c0f 100644 --- a/pkgs/development/cuda-modules/utils.nix +++ b/pkgs/development/cuda-modules/utils.nix @@ -1,20 +1,27 @@ -{ config, lib, ... }: +{ + config, + lib, + pkgs, + ... +}: let - inherit (lib.attrsets) - attrByPath - filterAttrs - mapAttrs - mapAttrs' - ; + inherit (lib.attrsets) filterAttrs mapAttrs mapAttrs'; inherit (lib.filesystem) packagesFromDirectoryRecursive; + inherit (lib.licenses) nvidiaCudaRedist; inherit (lib.lists) findFirst optionals reverseList take ; + inherit (lib.meta) addMetaAttrs; inherit (lib.options) mkOption; - inherit (lib.strings) concatStringsSep hasPrefix replaceStrings; + inherit (lib.strings) + concatStringsSep + hasPrefix + replaceStrings + versionAtLeast + ; inherit (lib.trivial) const flip @@ -29,12 +36,75 @@ let nonEmptyListOf nonEmptyStr nullOr - raw + package ; inherit (lib.versions) major majorMinor splitVersion; + inherit (pkgs) fetchurl srcOnly; in { options.utils = mapAttrs (const mkOption) { + buildRedistPackage = { + description = '' + Helper function which wraps the manifest builder to build a redistributable package. + Expects the fixed-point (`final`), and a flattenedIndexElem. + ''; + type = functionTo (functionTo package); + default = + final: + meta@{ + packageName, + redistName, + releaseInfo, + packageInfo, + version, + ... + }: + let + manifestBuilderArgs = { + inherit (meta) packageInfo packageName releaseInfo; + libPath = config.utils.getLibPath final.cudaMajorMinorPatchVersion packageInfo.feature.cudaVersionsInLib; + # The source is given by the tarball, which we unpack and use as a FOD. + src = + let + tarball = fetchurl { + inherit (packageInfo) sha256; + url = + if redistName == "tensorrt" then + config.utils.mkTensorRTURL version packageInfo.relativePath + else + config.utils.mkRedistURL redistName (config.utils.mkRelativePath meta); + }; + unpacked = srcOnly { + __structuredAttrs = true; + strictDeps = true; + name = tarball.name + "-unpacked"; + src = tarball; + outputHashMode = "recursive"; + outputHash = packageInfo.narHash; + }; + in + unpacked; + }; + package = pipe manifestBuilderArgs [ + # Build the package. + (final.callPackage ./manifest-builder.nix) + # Update the package license + (addMetaAttrs { + license = nvidiaCudaRedist // { + url = + let + licensePath = + if releaseInfo.licensePath != null then releaseInfo.licensePath else "${packageName}/LICENSE.txt"; + in + "https://developer.download.nvidia.com/compute/${redistName}/redist/${licensePath}"; + }; + }) + # Apply package-specific overrides, if they exist. + (config.utils.overridePackage final meta) + ]; + in + package; + }; getLibPath = { description = '' Returns the path to the CUDA library directory for a given version or null if no such version exists. @@ -59,52 +129,35 @@ in ] ); }; - getNixPlatform = { + getNixPlatforms = { description = '' - Function to map NVIDIA redist arch to Nix platform + Function to map NVIDIA redist arch to Nix platforms. + + NOTE: This function returns a list of platforms because the redistributable architecture `"source"` can be + built on multiple platforms. NOTE: This function *will* be called by unsupported platforms because `cudaPackages` is part of `all-packages.nix`, which is evaluated on all platforms. As such, we need to handle unsupported platforms gracefully. ''; - type = functionTo (enum [ - "aarch64-linux" - "ppc64le-linux" - "unsupported" - "windows-x86_64" - "x86_64-linux" - ]); + type = functionTo (nonEmptyListOf nonEmptyStr); default = - redistArch: - attrByPath [ redistArch ] "unsupported" { - linux-sbsa = "aarch64-linux"; - linux-aarch64 = "aarch64-linux"; - linux-x86_64 = "x86_64-linux"; - linux-ppc64le = "ppc64le-linux"; - windows-x86_64 = "x86_64-windows"; - }; - }; - getOverrides = { - description = "Return a nested attribute set of overrides for packages"; - type = - # let - # TODO: Can't typecheck this without getting - # error: cannot convert a function to JSON - # When using the overrideAttrsConst type (like cuda_nvml_dev). - # overrideAttrsPrevFn = functionTo overrideAttrsArg; - # overrideAttrsFinalPrevFn = functionTo overrideAttrsPrevFn; - # overrideAttrsArg = oneOf [ - # overrideAttrsPrevFn - # overrideAttrsFinalPrevFn - # ]; - # in - functionTo (config.types.attrs config.types.redistName (attrsOf raw)); - default = - final: - packagesFromDirectoryRecursive { - inherit (final) callPackage; - directory = ./overrides; - }; + let + platformMapping = { + linux-sbsa = [ "aarch64-linux" ]; + linux-aarch64 = [ "aarch64-linux" ]; + linux-x86_64 = [ "x86_64-linux" ]; + linux-ppc64le = [ "ppc64le-linux" ]; + windows-x86_64 = [ "x86_64-windows" ]; + source = [ + "aarch64-linux" # Both SBSA (ARM servers) and Jetson + "ppc64le-linux" # POWER + "x86_64-linux" # x86_64 + "x86_64-windows" # Windows + ]; + }; + in + redistArch: platformMapping.${redistArch} or [ "unsupported" ]; }; getRedistArch = { description = '' @@ -114,6 +167,7 @@ in `linux-aarch64` redist (which is for Jetson devices) if we're building any Jetson devices. Since both are based on aarch64, we can only have one or the other, otherwise there's an ambiguity as to which should be used. + NOTE: This function *will* be called by unsupported systems because `cudaPackages` is part of `all-packages.nix`, which is evaluated on all systems. As such, we need to handle unsupported systems gracefully. @@ -127,13 +181,15 @@ in "windows-x86_64" ]); default = - nixSystem: - attrByPath [ nixSystem ] "unsupported" { - aarch64-linux = if config.data.jetsonTargets != [ ] then "linux-aarch64" else "linux-sbsa"; - x86_64-linux = "linux-x86_64"; - ppc64le-linux = "linux-ppc64le"; - x86_64-windows = "windows-x86_64"; - }; + let + platformMapping = { + aarch64-linux = if config.data.jetsonTargets != [ ] then "linux-aarch64" else "linux-sbsa"; + x86_64-linux = "linux-x86_64"; + ppc64le-linux = "linux-ppc64le"; + x86_64-windows = "windows-x86_64"; + }; + in + nixSystem: platformMapping.${nixSystem} or "unsupported"; }; majorMinorPatch = { description = "Function to extract the major, minor, and patch version from a string"; @@ -182,6 +238,20 @@ in )) ]; }; + # mkSupportedPlatformsLookup = { + # description = '' + # Helper function to generate a lookup table of supported platforms for different packages. + # Takes cudaMajorMinorVersion and an index of packageInfo. + # ''; + # type = functionTo (functionTo (nonEmptyListOf config.types.platform)); + # default = + # cudaMajorMinorPatchVersion: packageInfoIndex: + # let + # cudaIndex = packageInfoIndex.cuda.${cudaMajorMinorPatchVersion}; + # in + # # supportedPlatforms = supportedPlatforms or [ ]; + # supportedPlatforms; + # }; mkTrimmedFlattenedIndex = { description = '' Helper function to reduce the number of attribute sets we need to traverse to create all the package sets. @@ -192,12 +262,12 @@ in ''; type = functionTo (functionTo (nonEmptyListOf config.types.flattenedIndexElem)); default = - cudaMajorMinorPatchVersion: index: + cudaMajorMinorPatchVersion: packageInfoIndex: let - allButCudaFlattenedIndex = config.utils.mkFlattenedIndex (removeAttrs index [ "cuda" ]); - maybeCudaFlattenedIndex = optionals (index.cuda ? ${cudaMajorMinorPatchVersion}) ( + allButCudaFlattenedIndex = config.utils.mkFlattenedIndex (removeAttrs packageInfoIndex [ "cuda" ]); + maybeCudaFlattenedIndex = optionals (packageInfoIndex.cuda ? ${cudaMajorMinorPatchVersion}) ( config.utils.mkFlattenedIndex { - cuda.${cudaMajorMinorPatchVersion} = index.cuda.${cudaMajorMinorPatchVersion}; + cuda.${cudaMajorMinorPatchVersion} = packageInfoIndex.cuda.${cudaMajorMinorPatchVersion}; } ); in @@ -214,5 +284,94 @@ in (version: "${name}_${version}") ]; }; + overridePackage = { + description = '' + A function to override a package. + Expects the fixed-point (`final`), a flattenedIndexElem, and the package to override. + ''; + type = functionTo (functionTo (functionTo package)); + default = + final: + let + # Maps redist name to package name to override function to provide to overrideAttrs. + overrides = packagesFromDirectoryRecursive { + inherit (final) callPackage; + directory = ./overrides; + }; + in + { packageName, redistName, ... }: + package: package.overrideAttrs (overrides.${redistName}.${packageName} or { }); + }; + packageSupportsCudaVersion = { + description = '' + Function to determine if a package supports a given CUDA version. + This is used to filter out packages that don't support the CUDA version we're building. + Expects a cudaMajorMinorPatchVersion and a flattenedIndexElem. + ''; + type = functionTo (functionTo bool); + default = + # 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? + cudaMajorMinorPatchVersion: + { + cudaVariant, + packageInfo, + platform, + redistName, + version, + ... + }: + let + inherit (packageInfo.feature) cudaVersionsInLib; + + # 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 = + (config.utils.getLibPath cudaMajorMinorPatchVersion cudaVersionsInLib) != null; + + # There is a variant for the desired CUDA version. + isDesiredCudaVariant = cudaVariant == (config.utils.mkCudaVariant cudaMajorMinorPatchVersion); + + # Default value for packages which don't specify some policy. + default = isDesiredCudaVariant || hasSupportedCudaVersionInLib; + + # Mapping + supportSet = { + # 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. + 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. + 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. + cutensor = hasSupportedCudaVersionInLib; + }; + in + supportSet.${redistName} or default; + }; }; }