nixpkgs/pkgs/development/cuda-modules/packages/libnvshmem.nix

{
  _cuda,
  backendStdenv,
  buildPackages,
  cmake,
  cccl,
  cuda_cudart,
  cuda_nvcc,
  cuda_nvml_dev,
  cuda_nvrtc,
  cuda_nvtx,
  cudaAtLeast,
  cudaMajorMinorVersion,
  cudaNamePrefix,
  fetchFromGitHub,
  flags,
  gdrcopy,
  lib,
  libfabric,
  libnvjitlink,
  mpi,
  nccl,
  ninja,
  pmix,
  python3Packages,
  rdma-core,
  ucx,
  # passthru.updateScript
  gitUpdater,

  withGdrcopy ? true,
  withIbgda ? true,
  withLibfabric ? true,
  withMpi ? true,
  withNccl ? nccl.meta.available,
  withPmix ? true,
  withUcx ? true,
}:
let
  inherit (lib)
    cmakeBool
    cmakeFeature
    getBin
    getDev
    getExe
    getLib
    licenses
    maintainers
    optional
    optionals
    teams
    ;
in
backendStdenv.mkDerivation (finalAttrs: {
  __structuredAttrs = true;
  strictDeps = true;

  # NOTE: Depends on the CUDA package set, so use cudaNamePrefix.
  name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
  pname = "libnvshmem";
  version = "3.6.5-0";

  src = fetchFromGitHub {
    owner = "NVIDIA";
    repo = "nvshmem";
    tag = "v${finalAttrs.version}";
    hash = "sha256-2E3/WGbg6srT/e3ykK0qxTy1ZlJ9JGGLlergG0ITwTY=";
  };

  outputs = [ "out" ];

  nativeBuildInputs = [
    cuda_nvcc
    cmake
    ninja

    # NOTE: Python is required even if not building nvshmem4py:
    # https://github.com/NVIDIA/nvshmem/blob/131da55f643ac87c810ba0bc51d359258bf433a1/CMakeLists.txt#L173
    python3Packages.python
  ]
  ++ optionals withMpi [
    # NOTE: mpi is in nativeBuildInputs because it contains compilers and is only discoverable by CMake
    # when a nativeBuildInput.
    mpi
  ];

  # NOTE: Hardcoded standard versions mean CMake doesn't respect values we provide, so we need to patch the files.
  postPatch = ''
    for standardName in {CXX,CUDA}_STANDARD
    do
      while IFS= read -r cmakeFileToPatch
      do
        nixLog "patching $PWD/$cmakeFileToPatch to fix $standardName"
        substituteInPlace "$PWD/$cmakeFileToPatch" \
          --replace-fail \
            "$standardName 11" \
            "$standardName 17"
      done < <(grep --recursive --files-with-matches "$standardName 11")
    done
    unset -v cmakeFileToPatch
    unset -v standardName
  '';

  enableParallelBuilding = true;

  buildInputs = [
    cccl
    cuda_cudart
    cuda_nvml_dev
    cuda_nvrtc
    cuda_nvtx
    libnvjitlink
    rdma-core
  ]
  ++ optionals withLibfabric [
    libfabric
  ]
  ++ optionals withGdrcopy [
    gdrcopy
  ]
  ++ optionals withNccl [
    nccl
  ]
  ++ optionals withPmix [
    pmix
  ]
  ++ optionals withUcx [
    ucx
  ];

  # NOTE: This *must* be an environment variable NVIDIA saw fit to *configure and build CMake projects* while *inside*
  # a CMake build and didn't correctly thread arguments through, so the environment is the only way to get
  # configurations to the nested build.
  env.CUDA_HOME = (getBin cuda_nvcc).outPath;

  # https://docs.nvidia.com/nvshmem/release-notes-install-guide/install-guide/nvshmem-install-proc.html#other-distributions
  cmakeFlags = lib.concatLists [
    [
      (cmakeFeature "NVSHMEM_PREFIX" (placeholder "out"))

      (cmakeFeature "CUDA_HOME" (getBin cuda_nvcc).outPath)
      (cmakeFeature "CMAKE_CUDA_COMPILER" (getExe cuda_nvcc))

      (cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)

      (cmakeBool "NVSHMEM_BUILD_TESTS" true)
      (cmakeBool "NVSHMEM_BUILD_EXAMPLES" true)

      (cmakeBool "NVSHMEM_BUILD_DEB_PACKAGE" false)
      (cmakeBool "NVSHMEM_BUILD_RPM_PACKAGE" false)

      # TODO: Looks like a nightmare to package and depends on things we haven't packaged yet
      # https://github.com/NVIDIA/nvshmem/tree/131da55f643ac87c810ba0bc51d359258bf433a1/nvshmem4py
      (cmakeBool "NVSHMEM_BUILD_PYTHON_LIB" false)

      # NOTE: unsupported because it requires Clang
      (cmakeBool "NVSHMEM_BUILD_BITCODE_LIBRARY" false)
    ]

    [ (cmakeBool "NVSHMEM_USE_NCCL" withNccl) ]
    (optional withNccl (cmakeFeature "NCCL_HOME" (getDev nccl).outPath))

    [ (cmakeBool "NVSHMEM_USE_GDRCOPY" withGdrcopy) ]
    (optional withGdrcopy (cmakeFeature "GDRCOPY_HOME" (getDev gdrcopy).outPath))

    # NOTE: Make sure to use mpi from buildPackages to match the spliced version created through nativeBuildInputs.
    [ (cmakeBool "NVSHMEM_MPI_SUPPORT" withMpi) ]
    (optional withMpi (cmakeFeature "MPI_HOME" (getLib buildPackages.mpi).outPath))

    # TODO: Doesn't UCX need to be built with some argument when we want to use it with libnvshmem?
    [ (cmakeBool "NVSHMEM_UCX_SUPPORT" withUcx) ]
    (optional withUcx (cmakeFeature "UCX_HOME" (getDev ucx).outPath))

    [ (cmakeBool "NVSHMEM_LIBFABRIC_SUPPORT" withLibfabric) ]
    (optional withLibfabric (cmakeFeature "LIBFABRIC_HOME" (getDev libfabric).outPath))

    # NOTE: no corresponding _HOME variable for IBGDA.
    [ (cmakeBool "NVSHMEM_IBGDA_SUPPORT" withIbgda) ]

    [ (cmakeBool "NVSHMEM_PMIX_SUPPORT" withPmix) ]
    (optional withPmix (cmakeFeature "PMIX_HOME" (getDev pmix).outPath))
  ];

  postInstall = ''
    nixLog "moving top-level files in $out to $out/share"
    mv -v "$out"/{changelog,git_commit.txt,License.txt,version.txt} "$out/share/"
  '';

  doCheck = false;

  passthru = {
    updateScript = gitUpdater {
      inherit (finalAttrs) pname version;
      rev-prefix = "v";
    };

    brokenAssertions = [
      # CUDA pre-11.7 yeilds macro/type errors in src/include/internal/host_transport/cudawrap.h.
      {
        message = "NVSHMEM does not support CUDA releases earlier than 11.7 (found ${cudaMajorMinorVersion})";
        assertion = cudaAtLeast "11.7";
      }
    ];
  };

  meta = {
    description = "Parallel programming interface for NVIDIA GPUs based on OpenSHMEM";
    homepage = "https://github.com/NVIDIA/nvshmem";
    changelog = "https://github.com/NVIDIA/nvshmem/releases/tag/${finalAttrs.src.tag}";
    broken = _cuda.lib._mkMetaBroken finalAttrs;
    # NOTE: There are many licenses:
    # https://github.com/NVIDIA/nvshmem/blob/7dd48c9fd7aa2134264400802881269b7822bd2f/License.txt
    license = licenses.nvidiaCudaRedist;
    platforms = [
      "aarch64-linux"
      "x86_64-linux"
    ];
    maintainers = with maintainers; [
      connorbaker
      GaetanLepage
    ];
    teams = [ teams.cuda ];
  };
})