diff options
Diffstat (limited to 'pkgs/development/cuda-modules/nccl/default.nix')
-rw-r--r-- | pkgs/development/cuda-modules/nccl/default.nix | 112 |
1 files changed, 112 insertions, 0 deletions
diff --git a/pkgs/development/cuda-modules/nccl/default.nix b/pkgs/development/cuda-modules/nccl/default.nix new file mode 100644 index 0000000000000..c56d59cb42068 --- /dev/null +++ b/pkgs/development/cuda-modules/nccl/default.nix @@ -0,0 +1,112 @@ +# NOTE: Though NCCL is called within the cudaPackages package set, we avoid passing in +# the names of dependencies from that package set directly to avoid evaluation errors +# in the case redistributable packages are not available. +{ + lib, + fetchFromGitHub, + python3, + which, + cudaPackages, + # passthru.updateScript + gitUpdater, +}: +let + inherit (cudaPackages) + autoAddOpenGLRunpathHook + backendStdenv + cuda_cccl + cuda_cudart + cuda_nvcc + cudaFlags + cudatoolkit + cudaVersion + ; +in +backendStdenv.mkDerivation ( + finalAttrs: { + pname = "nccl"; + version = "2.19.3-1"; + + src = fetchFromGitHub { + owner = "NVIDIA"; + repo = finalAttrs.pname; + rev = "v${finalAttrs.version}"; + hash = "sha256-59FlOKM5EB5Vkm4dZBRCkn+IgIcdQehE+FyZAdTCT/A="; + }; + + strictDeps = true; + + outputs = [ + "out" + "dev" + ]; + + nativeBuildInputs = + [ + which + autoAddOpenGLRunpathHook + python3 + ] + ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit] + ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [cuda_nvcc]; + + buildInputs = + lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit] + ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [ + cuda_nvcc.dev # crt/host_config.h + cuda_cudart + ] + # NOTE: CUDA versions in Nixpkgs only use a major and minor version. When we do comparisons + # against other version, like below, it's important that we use the same format. Otherwise, + # we'll get incorrect results. + # For example, lib.versionAtLeast "12.0" "12.0.0" == false. + ++ lib.optionals (lib.versionAtLeast cudaVersion "12.0") [cuda_cccl]; + + env.NIX_CFLAGS_COMPILE = toString ["-Wno-unused-function"]; + + preConfigure = '' + patchShebangs ./src/device/generate.py + makeFlagsArray+=( + "NVCC_GENCODE=${lib.concatStringsSep " " cudaFlags.gencode}" + ) + ''; + + makeFlags = + ["PREFIX=$(out)"] + ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [ + "CUDA_HOME=${cudatoolkit}" + "CUDA_LIB=${lib.getLib cudatoolkit}/lib" + "CUDA_INC=${lib.getDev cudatoolkit}/include" + ] + ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [ + "CUDA_HOME=${cuda_nvcc}" + "CUDA_LIB=${lib.getLib cuda_cudart}/lib" + "CUDA_INC=${lib.getDev cuda_cudart}/include" + ]; + + enableParallelBuilding = true; + + postFixup = '' + moveToOutput lib/libnccl_static.a $dev + ''; + + passthru.updateScript = gitUpdater { + inherit (finalAttrs) pname version; + rev-prefix = "v"; + }; + + meta = with lib; { + description = "Multi-GPU and multi-node collective communication primitives for NVIDIA GPUs"; + homepage = "https://developer.nvidia.com/nccl"; + license = licenses.bsd3; + platforms = platforms.linux; + maintainers = + with maintainers; + [ + mdaiter + orivej + ] + ++ teams.cuda.members; + }; + } +) |