Skip to content

Commit

Permalink
CUDA packages should be visible on all platforms, but not necessarily…
Browse files Browse the repository at this point in the history 
… buildable
  • Loading branch information
@ConnorBaker
ConnorBaker committed May 4, 2024
1 parent b340a8a commit 5732c52
Show file tree
Hide file tree
Showing 2 changed files with 155 additions and 86 deletions.
8 changes: 2 additions & 6 deletions pkgs/development/cuda-modules/manifest-builder.nix
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Builder-specific arguments
{
# Builder-specific arguments
# src :: Derivation
src,
# libPath :: null | Path
Expand All @@ -11,11 +11,7 @@
packageName,
# releaseInfo :: ReleaseInfo
releaseInfo,
# Various other args
...
}:
# General callPackage-supplied arguments
{
# General callPackage-supplied arguments
autoAddCudaCompatRunpath,
autoAddDriverRunpath,
autoPatchelfHook,
Expand Down
233 changes: 153 additions & 80 deletions pkgs/development/cuda-modules/package-sets.nix
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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";
Expand All @@ -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.
[
Expand All @@ -69,55 +64,112 @@ 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.
hasSupportedCudaVersionInLib =
(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,
...
Expand All @@ -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
Expand All @@ -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: {
Expand All @@ -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;
Expand Down

0 comments on commit 5732c52

Please sign in to comment.