cudaPackages: remove generic-builder

{

  # General callPackage-supplied arguments

  autoAddDriverRunpath,

  autoAddCudaCompatRunpath,

  autoPatchelfHook,

  backendStdenv,

  callPackage,

  _cuda,

  fetchurl,

  lib,

  markForCudatoolkitRootHook,

  flags,

  stdenv,

  # Builder-specific arguments

  # Short package name (e.g., "cuda_cccl")

  # pname : String

  pname,

  # Common name (e.g., "cutensor" or "cudnn") -- used in the URL.

  # Also known as the Redistributable Name.

  # redistName : String,

  redistName,

  # If libPath is non-null, it must be a subdirectory of `lib`.

  # The contents of `libPath` will be moved to the root of `lib`.

  libPath ? null,

  # See ./modules/generic/manifests/redistrib/release.nix

  redistribRelease,

  # See ./modules/generic/manifests/feature/release.nix

  featureRelease,

  cudaMajorMinorVersion,

}:

let

  inherit (lib)

    attrsets

    lists

    strings

    trivial

    licenses

    teams

    sourceTypes

    ;

  inherit (stdenv) hostPlatform;

  # Last step before returning control to `callPackage` (adds the `.override` method)

  # we'll apply (`overrideAttrs`) necessary package-specific "fixup" functions.

  # Order is significant.

  maybeFixup = _cuda.fixups.${pname} or null;

  fixup = if maybeFixup != null then callPackage maybeFixup { } else { };

  # Get the redist systems for which package provides distributables.

  # These are used by meta.platforms.

  supportedRedistSystems = builtins.attrNames featureRelease;

  # redistSystem :: String

  # The redistSystem is the name of the system for which the redistributable is built.

  # It is `"unsupported"` if the redistributable is not supported on the target system.

  redistSystem = _cuda.lib.getRedistSystem backendStdenv.hasJetsonCudaCapability hostPlatform.system;

  sourceMatchesHost = lib.elem hostPlatform.system (_cuda.lib.getNixSystems redistSystem);

in

(backendStdenv.mkDerivation (finalAttrs: {

  # NOTE: Even though there's no actual buildPhase going on here, the derivations of the

  # redistributables are sensitive to the compiler flags provided to stdenv. The patchelf package

  # is sensitive to the compiler flags provided to stdenv, and we depend on it. As such, we are

  # also sensitive to the compiler flags provided to stdenv.

  inherit pname;

  inherit (redistribRelease) version;

  # Don't force serialization to string for structured attributes, like outputToPatterns

  # and brokenConditions.

  # Avoids "set cannot be coerced to string" errors.

  __structuredAttrs = true;

  # Keep better track of dependencies.

  strictDeps = true;

  # NOTE: Outputs are evaluated jointly with meta, so in the case that this is an unsupported platform,

  # we still need to provide a list of outputs.

  outputs =

    let

      # Checks whether the redistributable provides an output.

      hasOutput =

        output:

        attrsets.attrByPath [

          redistSystem

          "outputs"

          output

        ] false featureRelease;

      # Order is important here so we use a list.

      possibleOutputs = [

        "bin"

        "lib"

        "static"

        "dev"

        "doc"

        "sample"

        "python"

      ];

      # Filter out outputs that don't exist in the redistributable.

      # NOTE: In the case the redistributable isn't supported on the target platform,

      # we will have `outputs = [ "out" ] ++ possibleOutputs`. This is of note because platforms which

      # aren't supported would otherwise have evaluation errors when trying to access outputs other than `out`.

      # The alternative would be to have `outputs = [ "out" ]` when`redistSystem = "unsupported"`, but that would

      # require adding guards throughout the entirety of the CUDA package set to ensure `cudaSupport` is true --

      # recall that OfBorg will evaluate packages marked as broken and that `cudaPackages` will be evaluated with

      # `cudaSupport = false`!

      additionalOutputs =

        if redistSystem == "unsupported" then

          possibleOutputs

        else

          builtins.filter hasOutput possibleOutputs;

      # The out output is special -- it's the default output and we always include it.

      outputs = [ "out" ] ++ additionalOutputs;

    in

    outputs;

  # Traversed in the order of the outputs specified in outputs;

  # entries are skipped if they don't exist in outputs.

  outputToPatterns = {

    bin = [ "bin" ];

    dev = [

      "share/pkgconfig"

      "**/*.pc"

      "**/*.cmake"

    ];

    lib = [

      "lib"

      "lib64"

    ];

    static = [ "**/*.a" ];

    sample = [ "samples" ];

    python = [ "**/*.whl" ];

  };

  # Useful for introspecting why something went wrong. Maps descriptions of why the derivation would be marked as

  # broken on have badPlatforms include the current platform.

  # brokenConditions :: AttrSet Bool

  # Sets `meta.broken = true` if any of the conditions are true.

  # Example: Broken on a specific version of CUDA or when a dependency has a specific version.

  brokenConditions = {

    # Unclear how this is handled by Nix internals.

    "Duplicate entries in outputs" = finalAttrs.outputs != lists.unique finalAttrs.outputs;

    # Typically this results in the static output being empty, as all libraries are moved

    # back to the lib output.

    "lib output follows static output" =

      let

        libIndex = lists.findFirstIndex (x: x == "lib") null finalAttrs.outputs;

        staticIndex = lists.findFirstIndex (x: x == "static") null finalAttrs.outputs;

      in

      libIndex != null && staticIndex != null && libIndex > staticIndex;

  };

  # badPlatformsConditions :: AttrSet Bool

  # Sets `meta.badPlatforms = meta.platforms` if any of the conditions are true.

  # Example: Broken on a specific architecture when some condition is met (like targeting Jetson).

  badPlatformsConditions = {

    "No source" = !sourceMatchesHost;

  };

  # src :: Optional Derivation

  # If redistSystem doesn't exist in redistribRelease, return null.

  src = trivial.mapNullable (

    { relative_path, sha256, ... }:

    fetchurl {

      url = "https://developer.download.nvidia.com/compute/${redistName}/redist/${relative_path}";

      inherit sha256;

    }

  ) (redistribRelease.${redistSystem} or null);

  postPatch =

    # Pkg-config's setup hook expects configuration files in $out/share/pkgconfig

    ''

      for path in pkg-config pkgconfig; do

        [[ -d "$path" ]] || continue

        mkdir -p share/pkgconfig

        mv "$path"/* share/pkgconfig/

        rmdir "$path"

      done

    ''

    # Rewrite FHS paths with store paths

    # NOTE: output* fall back to out if the corresponding output isn't defined.

    + ''

      for pc in share/pkgconfig/*.pc; do

        sed -i \

          -e "s|^cudaroot\s*=.*\$|cudaroot=''${!outputDev}|" \

          -e "s|^libdir\s*=.*/lib\$|libdir=''${!outputLib}/lib|" \

          -e "s|^includedir\s*=.*/include\$|includedir=''${!outputDev}/include|" \

          "$pc"

      done

    ''

    # Generate unversioned names.

    # E.g. cuda-11.8.pc -> cuda.pc

    + ''

      for pc in share/pkgconfig/*-"$majorMinorVersion.pc"; do

        ln -s "$(basename "$pc")" "''${pc%-$majorMinorVersion.pc}".pc

      done

    '';

  env.majorMinorVersion = cudaMajorMinorVersion;

  # We do need some other phases, like configurePhase, so the multiple-output setup hook works.

  dontBuild = true;

  nativeBuildInputs = [

    autoPatchelfHook

    # This hook will make sure libcuda can be found

    # in typically /lib/opengl-driver by adding that

    # directory to the rpath of all ELF binaries.

    # Check e.g. with `patchelf --print-rpath path/to/my/binary

    autoAddDriverRunpath

    markForCudatoolkitRootHook

  ]

  # autoAddCudaCompatRunpath depends on cuda_compat and would cause

  # infinite recursion if applied to `cuda_compat` itself (beside the fact

  # that it doesn't make sense in the first place)

  ++ lib.optionals (pname != "cuda_compat" && flags.isJetsonBuild) [

    # autoAddCudaCompatRunpath must appear AFTER autoAddDriverRunpath.

    # See its documentation in ./setup-hooks/extension.nix.

    autoAddCudaCompatRunpath

  ];

  buildInputs = [

    # autoPatchelfHook will search for a libstdc++ and we're giving it

    # one that is compatible with the rest of nixpkgs, even when

    # nvcc forces us to use an older gcc

    # NB: We don't actually know if this is the right thing to do

    (lib.getLib stdenv.cc.cc)

  ];

  # Picked up by autoPatchelf

  # Needed e.g. for libnvrtc to locate (dlopen) libnvrtc-builtins

  appendRunpaths = [ "$ORIGIN" ];

  # NOTE: We don't need to check for dev or doc, because those outputs are handled by

  # the multiple-outputs setup hook.

  # NOTE: moveToOutput operates on all outputs:

  # https://github.com/NixOS/nixpkgs/blob/2920b6fc16a9ed5d51429e94238b28306ceda79e/pkgs/build-support/setup-hooks/multiple-outputs.sh#L105-L107

  installPhase =

    let

      mkMoveToOutputCommand =

        output:

        let

          template = pattern: ''moveToOutput "${pattern}" "${"$" + output}"'';

          patterns = finalAttrs.outputToPatterns.${output} or [ ];

        in

        strings.concatMapStringsSep "\n" template patterns;

    in

    # Pre-install hook

    ''

      runHook preInstall

    ''

    # Handle the existence of libPath, which requires us to re-arrange the lib directory

    + strings.optionalString (libPath != null) ''

      full_lib_path="lib/${libPath}"

      if [[ ! -d "$full_lib_path" ]]; then

        echo "${finalAttrs.pname}: '$full_lib_path' does not exist, only found:" >&2

        find lib/ -mindepth 1 -maxdepth 1 >&2

        echo "This release might not support your CUDA version" >&2

        exit 1

      fi

      echo "Making libPath '$full_lib_path' the root of lib" >&2

      mv "$full_lib_path" lib_new

      rm -r lib

      mv lib_new lib

    ''

    # Create the primary output, out, and move the other outputs into it.

    + ''

      mkdir -p "$out"

      mv * "$out"

    ''

    # Move the outputs into their respective outputs.

    + strings.concatMapStringsSep "\n" mkMoveToOutputCommand (builtins.tail finalAttrs.outputs)

    # Add a newline to the end of the installPhase, so that the post-install hook doesn't

    # get concatenated with the last moveToOutput command.

    + "\n"

    # Post-install hook

    + ''

      runHook postInstall

    '';

  doInstallCheck = true;

  allowFHSReferences = true; # TODO: Default to `false`

  postInstallCheck = ''

    echo "Executing postInstallCheck"

    if [[ -z "''${allowFHSReferences-}" ]]; then

      mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "''${!o}"; done)

      if grep --max-count=5 --recursive --exclude=LICENSE /usr/ "''${outputPaths[@]}"; then

        echo "Detected references to /usr" >&2

        exit 1

      fi

    fi

  '';

  # libcuda needs to be resolved during runtime

  autoPatchelfIgnoreMissingDeps = [

    "libcuda.so"

    "libcuda.so.*"

  ];

  # _multioutPropagateDev() currently expects a space-separated string rather than an array

  preFixup = ''

    export propagatedBuildOutputs="''${propagatedBuildOutputs[@]}"

  '';

  # Propagate all outputs, including `static`

  propagatedBuildOutputs = builtins.filter (x: x != "dev") finalAttrs.outputs;

  # Kept in case overrides assume postPhases have already been defined

  postPhases = [ "postPatchelf" ];

  postPatchelf = ''

    true

  '';

  passthru = {

    # Provide access to the release information for fixup functions.

    inherit redistribRelease featureRelease;

    # Make the CUDA-patched stdenv available

    stdenv = backendStdenv;

  };

  meta = {

    description = "${redistribRelease.name}. By downloading and using the packages you accept the terms and conditions of the ${finalAttrs.meta.license.shortName}";

    sourceProvenance = [ sourceTypes.binaryNativeCode ];

    broken = lists.any trivial.id (attrsets.attrValues finalAttrs.brokenConditions);

    platforms = trivial.pipe supportedRedistSystems [

      # Map each redist system to the equivalent nix systems.

      (lib.concatMap _cuda.lib.getNixSystems)

      # Take all the unique values.

      lib.unique

      # Sort the list.

      lib.naturalSort

    ];

    badPlatforms =

      let

        isBadPlatform = lists.any trivial.id (attrsets.attrValues finalAttrs.badPlatformsConditions);

      in

      lists.optionals isBadPlatform finalAttrs.meta.platforms;

    license =

      if redistName == "cuda" then

        # Add the package-specific license.

        let

          licensePath =

            if redistribRelease.license_path != null then

              redistribRelease.license_path

            else

              "${pname}/LICENSE.txt";

          url = "https://developer.download.nvidia.com/compute/cuda/redist/${licensePath}";

        in

        lib.licenses.nvidiaCudaRedist // { inherit url; }

      else

        licenses.unfree;

    teams = [ teams.cuda ];

  };

})).overrideAttrs

  fixup

{

  lib,

  cudaLib,

  cudaMajorMinorVersion,

  redistSystem,

  stdenv,

  # Builder-specific arguments

  # Short package name (e.g., "cuda_cccl")

  # pname : String

  pname,

  # Common name (e.g., "cutensor" or "cudnn") -- used in the URL.

  # Also known as the Redistributable Name.

  # redistName : String,

  redistName,

  # releasesModule :: Path

  # A path to a module which provides a `releases` attribute

  releasesModule,

  # shims :: Path

  # A path to a module which provides a `shims` attribute

  # The redistribRelease is only used in ./manifest.nix for the package version

  # and the package description (which NVIDIA's manifest calls the "name").

  # It's also used for fetching the source, but we override that since we can't

  # re-use that portion of the functionality (different URLs, etc.).

  # The featureRelease is used to populate meta.platforms (by way of looking at the attribute names), determine the

  # outputs of the package, and provide additional package-specific constraints (e.g., min/max supported CUDA versions,

  # required versions of other packages, etc.).

  # shimFn :: {package, redistSystem} -> AttrSet

  shimsFn ? (throw "shimsFn must be provided"),

}:

let

  evaluatedModules = lib.modules.evalModules {

    modules = [

      ../modules

      releasesModule

    ];

  };

  # NOTE: Important types:

  # - Releases: ../modules/${pname}/releases/releases.nix

  # - Package: ../modules/${pname}/releases/package.nix

  # Check whether a package supports our CUDA version.

  # satisfiesCudaVersion :: Package -> Bool

  satisfiesCudaVersion =

    package:

    lib.versionAtLeast cudaMajorMinorVersion package.minCudaVersion

    && lib.versionAtLeast package.maxCudaVersion cudaMajorMinorVersion;

  # FIXME: do this at the module system level

  propagatePlatforms = lib.mapAttrs (redistSystem: lib.map (p: { inherit redistSystem; } // p));

  # Releases for all platforms and all CUDA versions.

  allReleases = propagatePlatforms evaluatedModules.config.${pname}.releases;

  # Releases for all platforms and our CUDA version.

  allReleases' = lib.mapAttrs (_: lib.filter satisfiesCudaVersion) allReleases;

  # Packages for all platforms and our CUDA versions.

  allPackages = lib.concatLists (lib.attrValues allReleases');

  packageOlder = p1: p2: lib.versionOlder p1.version p2.version;

  packageSupportedPlatform = p: p.redistSystem == redistSystem;

  # 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 :: Package -> String

  computeName = { version, ... }: cudaLib.mkVersionedName pname (lib.versions.majorMinor version);

  # The newest package for each major-minor version, with newest first.

  # newestPackages :: List Package

  newestPackages =

    let

      newestForEachMajorMinorVersion = lib.foldl' (

        newestPackages: package:

        let

          majorMinorVersion = lib.versions.majorMinor package.version;

          existingPackage = newestPackages.${majorMinorVersion} or null;

        in

        newestPackages

        // {

          ${majorMinorVersion} =

            # Only keep the existing package if it is newer than the one we are considering or it is supported on the

            # current platform and the one we are considering is not.

            if

              existingPackage != null

              && (

                packageOlder package existingPackage

                || (!packageSupportedPlatform package && packageSupportedPlatform existingPackage)

              )

            then

              existingPackage

            else

              package;

        }

      ) { } allPackages;

    in

    # Sort the packages by version so the newest is first.

    # NOTE: builtins.sort requires a strict weak ordering, so we must use versionOlder rather than versionAtLeast.

    # See https://github.com/NixOS/nixpkgs/commit/9fd753ea84e5035b357a275324e7fd7ccfb1fc77.

    lib.sort (lib.flip packageOlder) (lib.attrValues newestForEachMajorMinorVersion);

  extension =

    final: _:

    let

      # Builds our package into derivation and wraps it in a nameValuePair, where the name is the versioned name

      # of the package.

      buildPackage =

        package:

        let

          shims = final.callPackage shimsFn { inherit package redistSystem; };

          name = computeName package;

          drv = final.callPackage ./manifest.nix {

            inherit pname redistName;

            inherit (shims) redistribRelease featureRelease;

          };

        in

        lib.nameValuePair name drv;

      # versionedDerivations :: AttrSet Derivation

      versionedDerivations = builtins.listToAttrs (lib.map buildPackage newestPackages);

      defaultDerivation = {

        ${pname} = (buildPackage (lib.head newestPackages)).value;

      };

    in

    # NOTE: Must condition on the length of newestPackages to avoid non-total function lib.head aborting if

    # newestPackages is empty.

    lib.optionalAttrs (lib.length newestPackages > 0) (versionedDerivations // defaultDerivation);

in

extension