switch to utils.mkMissingPackagesBrokenConditions pattern to avoid mi…

…ssing attribute errors

pkgs/development/cuda-modules/overrides/cublasmp/libcublasmp.nix

@@ -1,5 +1,13 @@
-{ libcublas, libcal }:
+args@{
+  lib,
+  libcal ? null,
+  libcublas ? null,
+  utils,
+}:
 prevAttrs: {
+  brokenConditions = prevAttrs.brokenConditions // (utils.mkMissingPackagesBrokenConditions args);
+  # TODO: Looks like the minimum supported capability is 7.0 as of the latest:
+  # https://docs.nvidia.com/cuda/cublasmp/getting_started/index.html
   buildInputs = prevAttrs.buildInputs or [ ] ++ [
     libcal.lib
     libcublas.lib

pkgs/development/cuda-modules/overrides/cuquantum/cuquantum.nix

@@ -1,12 +1,16 @@
-{
-  libcutensor,
-  libcublas,
-  libcusolver,
+args@{
+  cudaOlder,
+  lib,
+  libcublas ? null,
+  libcusolver ? null,
+  libcutensor ? null,
+  utils,
 }:
 prevAttrs: {
+  brokenConditions = prevAttrs.brokenConditions // (utils.mkMissingPackagesBrokenConditions args);
   buildInputs = prevAttrs.buildInputs or [ ] ++ [
-    libcutensor.lib
     libcublas.lib
     libcusolver.lib
+    libcutensor.lib
   ];
 }

pkgs/development/cuda-modules/overrides/cusolvermp/libcusolvermp.nix

@@ -1,10 +1,13 @@
-{
-  libcublas,
-  libcal,
-  libcusolver,
-  cuda_cudart,
+args@{
+  cuda_cudart ? null,
+  lib,
+  libcal ? null,
+  libcublas ? null,
+  libcusolver ? null,
+  utils,
 }:
 prevAttrs: {
+  brokenConditions = prevAttrs.brokenConditions // (utils.mkMissingPackagesBrokenConditions args);
   buildInputs = prevAttrs.buildInputs or [ ] ++ [
     libcal.lib
     libcublas.lib

pkgs/development/cuda-modules/package-sets.nix

@@ -13,7 +13,7 @@ let
     stdenv
     ;
   inherit (lib.attrsets)
-    attrNames
+    attrByPath
     dontRecurseIntoAttrs
     filterAttrs
     mapAttrs
@@ -134,31 +134,51 @@ let
       ...
     }:
     let
-      # 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.
+      # 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
-    if redistName == "cuda" then
+    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.
-      version == cudaMajorMinorPatchVersion
-    else if redistName == "cudnn" then
+      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.
-      isDesiredCudaVariant
-    else if redistName == "cutensor" then
+      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.
-      hasSupportedCudaVersionInLib
-    else
-      # TODO:
-      (isDesiredCudaVariant || hasSupportedCudaVersionInLib);
+      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 =

pkgs/development/cuda-modules/utils.nix

@@ -1,6 +1,11 @@
 { config, lib, ... }:
 let
-  inherit (lib.attrsets) attrByPath mapAttrs;
+  inherit (lib.attrsets)
+    attrByPath
+    filterAttrs
+    mapAttrs
+    mapAttrs'
+    ;
   inherit (lib.filesystem) packagesFromDirectoryRecursive;
   inherit (lib.lists)
     findFirst
@@ -18,6 +23,7 @@ let
     ;
   inherit (lib.types)
     attrsOf
+    bool
     enum
     functionTo
     nonEmptyListOf
@@ -159,6 +165,23 @@ in
         (builtins.removeAttrs args [ "leaf" ]) // { packageInfo = leaf; }
       );
     };
+    mkMissingPackagesBrokenConditions = {
+      description = ''
+        Utility function to generate a set of brokenConditions for missing packages.
+      '';
+      type = functionTo (attrsOf bool);
+      default = flip pipe [
+        # Take the attributes that are null.
+        (filterAttrs (_: value: value == null))
+        # Map them to a set of brokenConditions.
+        (mapAttrs' (
+          name: value: {
+            name = "Required package ${name} is missing";
+            value = true;
+          }
+        ))
+      ];
+    };
     mkTrimmedFlattenedIndex = {
       description = ''
         Helper function to reduce the number of attribute sets we need to traverse to create all the package sets.