646b892680
Some checks failed
Periodic Merges (6h) / master → staging-nixos (push) Failing after 12m50s
Periodic Merges (6h) / master → staging-next (push) Failing after 12m54s
Periodic Merges (24h) / merge-base(master,staging) → haskell-updates (push) Failing after 11m54s
Periodic Merges (6h) / staging-next → staging (push) Failing after 12m13s
Periodic Merges (24h) / staging-next-25.05 → staging-25.05 (push) Failing after 13m24s
Periodic Merges (24h) / release-25.05 → staging-next-25.05 (push) Failing after 14m28s
125 lines
4.0 KiB
Nix
125 lines
4.0 KiB
Nix
# Support matrix can be found at
|
|
# https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-880/support-matrix/index.html
|
|
#
|
|
# TODO(@connorbaker):
|
|
# This is a very similar strategy to CUDA/CUDNN:
|
|
#
|
|
# - Get all versions supported by the current release of CUDA
|
|
# - Build all of them
|
|
# - Make the newest the default
|
|
#
|
|
# Unique twists:
|
|
#
|
|
# - Instead of providing different releases for each version of CUDA, CuTensor has multiple subdirectories in `lib`
|
|
# -- one for each version of CUDA.
|
|
{
|
|
cudaLib,
|
|
cudaMajorMinorVersion,
|
|
lib,
|
|
redistSystem,
|
|
}:
|
|
let
|
|
inherit (lib)
|
|
attrsets
|
|
lists
|
|
modules
|
|
versions
|
|
trivial
|
|
;
|
|
|
|
redistName = "cutensor";
|
|
pname = "libcutensor";
|
|
|
|
cutensorVersions = [
|
|
"2.0.2"
|
|
"2.1.0"
|
|
];
|
|
|
|
# Manifests :: { redistrib, feature }
|
|
|
|
# Each release of cutensor gets mapped to an evaluated module for that release.
|
|
# From there, we can get the min/max CUDA versions supported by that release.
|
|
# listOfManifests :: List Manifests
|
|
listOfManifests =
|
|
let
|
|
configEvaluator =
|
|
fullCutensorVersion:
|
|
modules.evalModules {
|
|
modules = [
|
|
../modules
|
|
# We need to nest the manifests in a config.cutensor.manifests attribute so the
|
|
# module system can evaluate them.
|
|
{
|
|
cutensor.manifests = {
|
|
redistrib = trivial.importJSON (./manifests + "/redistrib_${fullCutensorVersion}.json");
|
|
feature = trivial.importJSON (./manifests + "/feature_${fullCutensorVersion}.json");
|
|
};
|
|
}
|
|
];
|
|
};
|
|
# Un-nest the manifests attribute set.
|
|
releaseGrabber = evaluatedModules: evaluatedModules.config.cutensor.manifests;
|
|
in
|
|
lists.map (trivial.flip trivial.pipe [
|
|
configEvaluator
|
|
releaseGrabber
|
|
]) cutensorVersions;
|
|
|
|
# Our cudaMajorMinorVersion tells us which version of CUDA we're building against.
|
|
# The subdirectories in lib/ tell us which versions of CUDA are supported.
|
|
# Typically the names will look like this:
|
|
#
|
|
# - 11
|
|
# - 12
|
|
|
|
# libPath :: String
|
|
libPath = versions.major cudaMajorMinorVersion;
|
|
|
|
# A release is supported if it has a libPath that matches our CUDA version for our platform.
|
|
# LibPath are not constant across the same release -- one platform may support fewer
|
|
# CUDA versions than another.
|
|
# platformIsSupported :: Manifests -> Boolean
|
|
platformIsSupported =
|
|
{ feature, redistrib, ... }:
|
|
(attrsets.attrByPath [
|
|
pname
|
|
redistSystem
|
|
] null feature) != null;
|
|
|
|
# TODO(@connorbaker): With an auxiliary file keeping track of the CUDA versions each release supports,
|
|
# we could filter out releases that don't support our CUDA version.
|
|
# However, we don't have that currently, so we make a best-effort to try to build TensorRT with whatever
|
|
# libPath corresponds to our CUDA version.
|
|
# supportedManifests :: List Manifests
|
|
supportedManifests = builtins.filter platformIsSupported listOfManifests;
|
|
|
|
# Compute versioned attribute name to be used in this package set
|
|
# Patch version changes should not break the build, so we only use major and minor
|
|
# computeName :: RedistribRelease -> String
|
|
computeName =
|
|
{ version, ... }: cudaLib.mkVersionedName redistName (lib.versions.majorMinor version);
|
|
in
|
|
final: _:
|
|
let
|
|
# buildCutensorPackage :: Manifests -> AttrSet Derivation
|
|
buildCutensorPackage =
|
|
{ redistrib, feature }:
|
|
let
|
|
drv = final.callPackage ../generic-builders/manifest.nix {
|
|
inherit pname redistName libPath;
|
|
redistribRelease = redistrib.${pname};
|
|
featureRelease = feature.${pname};
|
|
};
|
|
in
|
|
attrsets.nameValuePair (computeName redistrib.${pname}) drv;
|
|
|
|
extension =
|
|
let
|
|
nameOfNewest = computeName (lists.last supportedManifests).redistrib.${pname};
|
|
drvs = builtins.listToAttrs (lists.map buildCutensorPackage supportedManifests);
|
|
containsDefault = attrsets.optionalAttrs (drvs != { }) { cutensor = drvs.${nameOfNewest}; };
|
|
in
|
|
drvs // containsDefault;
|
|
in
|
|
extension
|