pyrox.dev
1# Support matrix can be found at
2# https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-880/support-matrix/index.html
3{
4 cudaLib,
5 lib,
6 redistSystem,
7}:
8let
9 inherit (lib)
10 attrsets
11 lists
12 modules
13 trivial
14 ;
15
16 redistName = "cusparselt";
17 pname = "libcusparse_lt";
18
19 cusparseltVersions = [
20 "0.7.1"
21 ];
22
23 # Manifests :: { redistrib, feature }
24
25 # Each release of cusparselt gets mapped to an evaluated module for that release.
26 # From there, we can get the min/max CUDA versions supported by that release.
27 # listOfManifests :: List Manifests
28 listOfManifests =
29 let
30 configEvaluator =
31 fullCusparseltVersion:
32 modules.evalModules {
33 modules = [
34 ../modules
35 # We need to nest the manifests in a config.cusparselt.manifests attribute so the
36 # module system can evaluate them.
37 {
38 cusparselt.manifests = {
39 redistrib = trivial.importJSON (./manifests + "/redistrib_${fullCusparseltVersion}.json");
40 feature = trivial.importJSON (./manifests + "/feature_${fullCusparseltVersion}.json");
41 };
42 }
43 ];
44 };
45 # Un-nest the manifests attribute set.
46 releaseGrabber = evaluatedModules: evaluatedModules.config.cusparselt.manifests;
47 in
48 lists.map (trivial.flip trivial.pipe [
49 configEvaluator
50 releaseGrabber
51 ]) cusparseltVersions;
52
53 # platformIsSupported :: Manifests -> Boolean
54 platformIsSupported =
55 { feature, redistrib, ... }:
56 (attrsets.attrByPath [
57 pname
58 redistSystem
59 ] null feature) != null;
60
61 # TODO(@connorbaker): With an auxiliary file keeping track of the CUDA versions each release supports,
62 # we could filter out releases that don't support our CUDA version.
63 # However, we don't have that currently, so we make a best-effort to try to build TensorRT with whatever
64 # libPath corresponds to our CUDA version.
65 # supportedManifests :: List Manifests
66 supportedManifests = builtins.filter platformIsSupported listOfManifests;
67
68 # Compute versioned attribute name to be used in this package set
69 # Patch version changes should not break the build, so we only use major and minor
70 # computeName :: RedistribRelease -> String
71 computeName =
72 { version, ... }: cudaLib.mkVersionedName redistName (lib.versions.majorMinor version);
73in
74final: _:
75let
76 # buildCusparseltPackage :: Manifests -> AttrSet Derivation
77 buildCusparseltPackage =
78 { redistrib, feature }:
79 let
80 drv = final.callPackage ../generic-builders/manifest.nix {
81 inherit pname redistName;
82 redistribRelease = redistrib.${pname};
83 featureRelease = feature.${pname};
84 };
85 in
86 attrsets.nameValuePair (computeName redistrib.${pname}) drv;
87
88 extension =
89 let
90 nameOfNewest = computeName (lists.last supportedManifests).redistrib.${pname};
91 drvs = builtins.listToAttrs (lists.map buildCusparseltPackage supportedManifests);
92 containsDefault = attrsets.optionalAttrs (drvs != { }) { cusparselt = drvs.${nameOfNewest}; };
93 in
94 drvs // containsDefault;
95in
96extension