From ef768dbd5058d939739ab63592133760c4850515 Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Mon, 12 May 2025 20:32:46 +0000 Subject: [PATCH] pkgsForCudaArch: init Signed-off-by: Connor Baker --- pkgs/top-level/variants.nix | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/pkgs/top-level/variants.nix b/pkgs/top-level/variants.nix index dfda26f60d08..ad4748f8b782 100644 --- a/pkgs/top-level/variants.nix +++ b/pkgs/top-level/variants.nix @@ -95,6 +95,27 @@ self: super: { }; }); + # `pkgsForCudaArch` maps each CUDA capability in _cuda.db.cudaCapabilityToInfo to a Nixpkgs variant configured for + # that target system. For example, `pkgsForCudaArch.sm_90a.python3Packages.torch` refers to PyTorch built for the + # Hopper architecture, leveraging architecture-specific features. + # NOTE: Not every package set is supported on every architecture! + # See `Using pkgsForCudaArch` in doc/languages-frameworks/cuda.section.md for more information. + pkgsForCudaArch = lib.listToAttrs ( + lib.map (cudaCapability: { + name = self._cuda.lib.mkRealArchitecture cudaCapability; + value = nixpkgsFun { + config = super.config // { + cudaSupport = true; + rocmSupport = false; + # Not supported by architecture-specific feature sets, so disable for all. + # Users can choose to build for family-specific feature sets if they wish. + cudaForwardCompat = false; + cudaCapabilities = [ cudaCapability ]; + }; + }; + }) (lib.attrNames self._cuda.db.cudaCapabilityToInfo) + ); + pkgsExtraHardening = nixpkgsFun { overlays = [ (