From 5732c52971e550cba894bf91518d9859009105f1 Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Fri, 3 May 2024 23:35:41 -0700 Subject: [PATCH] CUDA packages should be visible on all platforms, but not necessarily buildable --- .../cuda-modules/manifest-builder.nix | 8 +- .../development/cuda-modules/package-sets.nix | 233 ++++++++++++------ 2 files changed, 155 insertions(+), 86 deletions(-) diff --git a/pkgs/development/cuda-modules/manifest-builder.nix b/pkgs/development/cuda-modules/manifest-builder.nix index 706316ec012e978..454a54c3009990f 100644 --- a/pkgs/development/cuda-modules/manifest-builder.nix +++ b/pkgs/development/cuda-modules/manifest-builder.nix @@ -1,5 +1,5 @@ -# Builder-specific arguments { + # Builder-specific arguments # src :: Derivation src, # libPath :: null | Path @@ -11,11 +11,7 @@ packageName, # releaseInfo :: ReleaseInfo releaseInfo, - # Various other args - ... -}: -# General callPackage-supplied arguments -{ + # General callPackage-supplied arguments autoAddCudaCompatRunpath, autoAddDriverRunpath, autoPatchelfHook, diff --git a/pkgs/development/cuda-modules/package-sets.nix b/pkgs/development/cuda-modules/package-sets.nix index 0e6ac22d78ae934..0ec610caf613d7b 100644 --- a/pkgs/development/cuda-modules/package-sets.nix +++ b/pkgs/development/cuda-modules/package-sets.nix @@ -22,7 +22,13 @@ let inherit (lib.customisation) makeScope; inherit (lib.filesystem) packagesFromDirectoryRecursive; inherit (lib.licenses) nvidiaCudaRedist; - inherit (lib.lists) foldl' map; + inherit (lib.lists) + foldl' + map + naturalSort + optionals + unique + ; inherit (lib.meta) addMetaAttrs; inherit (lib.options) mkOption; inherit (lib.strings) removeSuffix versionAtLeast versionOlder; @@ -33,6 +39,7 @@ let # TODO: # - Version constraint handling (like for cutensor) # - 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"; @@ -47,20 +54,8 @@ let { packageName, redistName, ... }: package: package.overrideAttrs (overrides.${redistName}.${packageName} or { }); - # Function to get supported Nix systems for a package given a flattenedIndexElem. - getSupportedNixSystems = - cudaMajorMinorPatchVersion: - let - desiredCudaVariant = utils.mkCudaVariant cudaMajorMinorPatchVersion; - in - { - cudaVariant, - packageName, - platform, - redistName, - version, - ... - }: + getNixPlatforms = + platform: if platform == "source" then # All platforms are supported for source packages. [ @@ -69,27 +64,76 @@ let "ppc64le-linux" # POWER ] else - # Take all platforms which have a key with the desired CUDA variant or None. - pipe data.indices.packageInfo.${redistName}.${version}.${packageName} [ - (filterAttrs ( - _: cudaVariantAttrSet: cudaVariantAttrSet ? ${desiredCudaVariant} || cudaVariantAttrSet ? None - )) - attrNames - (map utils.getNixPlatform) + [ (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. - acceptPackage = + packageCudaVariantMatches = cudaMajorMinorPatchVersion: { cudaVariant, packageInfo, platform, redistName, + version, ... }: let - isSupportedPlatform = platform == "source" || (isSupportedRedistArch && platform == redistArch); # 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. @@ -97,27 +141,35 @@ let (utils.getLibPath cudaMajorMinorPatchVersion packageInfo.feature.cudaVersionsInLib) != null; in if redistName == "cuda" then - isSupportedPlatform + # 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 # 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. - isSupportedPlatform && isDesiredCudaVariant + isDesiredCudaVariant else if redistName == "cutensor" then # 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. - isSupportedPlatform && hasSupportedCudaVersionInLib + hasSupportedCudaVersionInLib else # TODO: - isSupportedPlatform && (isDesiredCudaVariant || hasSupportedCudaVersionInLib); + (isDesiredCudaVariant || hasSupportedCudaVersionInLib); + + # 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 = { + packageInfo, packageName, + platform, redistName, releaseInfo, ... @@ -132,11 +184,75 @@ let else packageName; - # Booleans which instruct us on how to update the package set at this step. - packagesHasEntry = packages ? ${attributeName}; - entryIsOlder = versionOlder packages.${attributeName}.version package.version; - packagesHasDefault = packages ? ${packageName}; - defaultIsOlder = versionOlder packages.${packageName}.version package.version; + # 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. + packageForName = + name: + let + replacePackage = + # There's no entry for the package in the package set. + !(packages ? ${name}) + # There is an entry but it's older than the package we're processing. + || (versionOlder packages.${name}.version package.version); + + package' = if replacePackage then package else packages.${name}; + + platformMatched = packagePlatformMatches platform; + in + pipe package' [ + # Update the package. + ( + pkg: + # Only do the override if the platform matched or we're replacing the package. + # Otherwise leave it alone. + pkg.override ( + optionalAttrs (platformMatched || replacePackage) { + src = + if platformMatched then + package.src + else if replacePackage then + null + else + builtins.throw "This should never happen."; + packageInfo = + if platformMatched then + packageInfo + else if replacePackage then + { feature.outputs = [ "out" ]; } + else + builtins.throw "This should never happen."; + } + ) + ) + # We need to append the current platform to the `meta.platforms` attribute. + ( + pkg: + pkg.overrideAttrs (prevAttrs: { + meta.platforms = + let + existingPlatforms = optionals (!replacePackage) (prevAttrs.meta.platforms or [ ]); + newPlatforms = 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... + unique (naturalSort (existingPlatforms ++ newPlatforms)); + }) + ) + ]; in packages # We keep point-releases for non-CUDA redist packages because NVIDIA may not release a new point release for @@ -146,12 +262,10 @@ let # NOTE: In the case we're processing a CUDA redistributable, the attribute name and the package name are the same, # so we're effectively replacing the package twice. // { - ${attributeName} = - if (!packagesHasEntry || entryIsOlder) then package else packages.${attributeName}; + ${attributeName} = packageForName attributeName; } // { - ${packageName} = - if (!packagesHasDefault || defaultIsOlder) then package else packages.${packageName}; + ${packageName} = packageForName packageName; }; packageSetBuilder = cudaMajorMinorPatchVersion: { @@ -170,50 +284,9 @@ let ... }: let - tarballSrc = fetchurl { - inherit (packageInfo) sha256; - url = - if redistName == "tensorrt" then - utils.mkTensorRTURL version packageInfo.relativePath - else - utils.mkRedistURL redistName (utils.mkRelativePath meta); - }; - src = srcOnly { - __structuredAttrs = true; - strictDeps = true; - - name = pipe tarballSrc.name [ - (removeSuffix ".tar.xz") - (removeSuffix ".tar.gz") - ]; - src = tarballSrc; - - outputHashMode = "recursive"; - outputHash = packageInfo.narHash; - }; - libPath = utils.getLibPath cudaMajorMinorPatchVersion packageInfo.feature.cudaVersionsInLib; - package = pipe ./manifest-builder.nix [ - # Import the builder function. - (flip builtins.import (meta // { inherit libPath src; })) - # Build the package with the final scope. - (flip final.callPackage { }) - # Update the package license and platforms - (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}"; - }; - platforms = getSupportedNixSystems cudaMajorMinorPatchVersion meta; - }) - # Apply package-specific overrides, if they exist. - (overridePackage final meta) - ]; + package = buildRedistPackage final cudaMajorMinorPatchVersion meta; in - if acceptPackage cudaMajorMinorPatchVersion meta then + if packageCudaVariantMatches cudaMajorMinorPatchVersion meta then updatePackages meta packages package else packages;