cudaPackages.cudnn-frontend: init at 1.8.0

pkgs/development/cuda-modules/cudnn-frontend/CMakeLists.txt

@@ -0,0 +1,133 @@
+cmake_minimum_required(VERSION 3.23)
+
+project(cudnn_frontend VERSION 1.8.0)
+
+option(CUDNN_FRONTEND_SKIP_JSON_LIB "Defines whether FE should not include nlohmann/json.hpp." OFF)
+option(CUDNN_FRONTEND_BUILD_SAMPLES "Defines if samples are built or not." ON)
+option(CUDNN_FRONTEND_BUILD_TESTS "Defines if unittests are built or not." ON)
+option(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS "Defines if python bindings are built or not." OFF)
+
+if(MSVC OR MSYS OR MINGW)
+    add_compile_options(/W4 /WX)
+else()
+    add_compile_options(-Wall -Wextra -Wpedantic -Werror -Wno-error=attributes -Wno-attributes -Wno-error=unused-function -Wno-unused-function)
+endif()
+
+add_library(cudnn_frontend INTERFACE)
+
+# Add header files to library
+file(GLOB_RECURSE CUDNN_FRONTEND_INCLUDE_FILES "include/*")
+target_sources(
+    cudnn_frontend
+    PUBLIC 
+        FILE_SET
+            HEADERS
+            BASE_DIRS
+                "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
+            FILES
+                "${CUDNN_FRONTEND_INCLUDE_FILES}"
+)
+unset(CUDNN_FRONTEND_INCLUDE_FILES)
+
+target_compile_definitions(cudnn_frontend INTERFACE $<$<BOOL:${CUDNN_FRONTEND_SKIP_JSON_LIB}>:CUDNN_FRONTEND_SKIP_JSON_LIB>)
+
+target_include_directories(
+    cudnn_frontend
+    INTERFACE
+        "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
+        "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>"
+)
+
+# Find the cuda compiler
+find_package(CUDAToolkit REQUIRED)
+
+target_include_directories(cudnn_frontend INTERFACE ${CUDAToolkit_INCLUDE_DIRS})
+
+target_compile_features(cudnn_frontend INTERFACE cxx_std_17)
+
+# Make PCH for targets to link against
+add_library(_cudnn_frontend_pch INTERFACE)
+target_precompile_headers(_cudnn_frontend_pch INTERFACE ${PROJECT_SOURCE_DIR}/include/cudnn_frontend.h)
+
+if (CUDNN_FRONTEND_BUILD_SAMPLES)
+    add_subdirectory(samples)
+    target_link_libraries(
+        samples 
+        PRIVATE
+            CUDA::cublasLt
+            CUDA::nvrtc
+    )
+    target_link_libraries(
+        legacy_samples
+        PRIVATE
+            CUDA::cublasLt
+            CUDA::nvrtc
+    )
+endif()
+
+if (CUDNN_FRONTEND_BUILD_TESTS)
+    add_subdirectory(test)
+    target_link_libraries(
+        tests
+        CUDA::cublasLt
+        CUDA::nvrtc
+    )
+endif()
+
+if (CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS)
+    add_subdirectory(python)
+endif()
+
+# Introduce variables:
+# * CMAKE_INSTALL_LIBDIR
+# * CMAKE_INSTALL_BINDIR
+# * CMAKE_INSTALL_INCLUDEDIR
+include(GNUInstallDirs)
+
+# Install and export the header files
+install(
+    TARGETS 
+        cudnn_frontend
+    EXPORT
+        cudnn_frontend_targets
+    FILE_SET HEADERS
+)
+
+if (CUDNN_FRONTEND_BUILD_SAMPLES)
+    install(TARGETS legacy_samples samples RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+endif()
+
+if (CUDNN_FRONTEND_BUILD_TESTS)
+    install(TARGETS tests RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+endif()
+
+# See https://cmake.org/cmake/help/latest/module/CMakePackageConfigHelpers.html#example-generating-package-files
+include(CMakePackageConfigHelpers)
+
+export(
+    EXPORT
+        cudnn_frontend_targets
+    FILE
+        "${CMAKE_CURRENT_BINARY_DIR}/cudnn_frontend/cudnn_frontend-targets.cmake"
+)
+install(
+    EXPORT
+        cudnn_frontend_targets
+    FILE
+        cudnn_frontend-targets.cmake
+    DESTINATION
+        "${CMAKE_INSTALL_LIBDIR}/cmake/cudnn_frontend"
+)
+
+configure_package_config_file(
+    cudnn_frontend-config.cmake.in
+    "${CMAKE_CURRENT_BINARY_DIR}/cudnn_frontend-config.cmake"
+    INSTALL_DESTINATION
+        "${CMAKE_INSTALL_LIBDIR}/cmake/cudnn_frontend"
+)
+install(
+    FILES
+        "${CMAKE_CURRENT_BINARY_DIR}/cudnn_frontend-config.cmake"
+    DESTINATION
+        "${CMAKE_INSTALL_LIBDIR}/cmake/cudnn_frontend"
+)

pkgs/development/cuda-modules/cudnn-frontend/cudnn_frontend-config.cmake.in

@@ -0,0 +1,3 @@
+@PACKAGE_INIT@
+
+include(${CMAKE_CURRENT_LIST_DIR}/cudnn_frontend-targets.cmake)

pkgs/development/cuda-modules/cudnn-frontend/default.nix

@@ -0,0 +1,131 @@
+{
+  autoAddDriverRunpath,
+  catch2_3,
+  cmake,
+  cuda_cccl,
+  cuda_cudart,
+  cuda_nvcc,
+  cuda_nvrtc,
+  cudnn,
+  fetchFromGitHub,
+  gitUpdater,
+  lib,
+  libcublas,
+  ninja,
+  nlohmann_json,
+  stdenv,
+}:
+let
+  inherit (lib.lists) optionals;
+  inherit (lib.strings)
+    cmakeBool
+    cmakeFeature
+    optionalString
+    ;
+in
+
+# TODO(@connorbaker): This should be a hybrid C++/Python package.
+stdenv.mkDerivation (finalAttrs: {
+  pname = "cudnn-frontend";
+  version = "1.8.0";
+
+  src = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cudnn-frontend";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-hKqIWGxVco1qkKxDZjc+pUisIcYJwFjZobJZg1WgDvY=";
+  };
+
+  # TODO: As a header-only library, we should make sure we have an `include` directory or similar which is not a
+  # superset of the `out` (`bin`) or `dev` outputs (whih is what the multiple-outputs setup hook does by default).
+  outputs =
+    [
+      "out"
+    ]
+    ++ optionals finalAttrs.doCheck [
+      "legacy_samples"
+      "samples"
+      "tests"
+    ];
+
+  nativeBuildInputs = [
+    autoAddDriverRunpath # Needed for samples because it links against CUDA::cuda_driver
+    cmake
+    cuda_cccl
+    cuda_nvcc
+    ninja
+  ];
+
+  buildInputs = [
+    cuda_cudart
+  ];
+
+  postPatch =
+    # nlohmann_json should be the only vendored dependency.
+    ''
+      echo "patching source to use nlohmann_json from nixpkgs"
+      rm -rf include/cudnn_frontend/thirdparty/nlohmann
+      rmdir include/cudnn_frontend/thirdparty
+      substituteInPlace include/cudnn_frontend_utils.h \
+        --replace-fail \
+          '#include "cudnn_frontend/thirdparty/nlohmann/json.hpp"' \
+          '#include <nlohmann/json.hpp>'
+    ''
+    # Use our own CMakeLists.txt and add a config to make the include files discoverable
+    + ''
+      echo "replacing CMakeLists.txt to link against forgotten libraries and install targets"
+      rm ./CMakeLists.txt
+      install -Dm644 "${./CMakeLists.txt}" ./CMakeLists.txt
+      install -Dm644 "${./cudnn_frontend-config.cmake.in}" ./cudnn_frontend-config.cmake.in
+    '';
+
+  cmakeFlags = [
+    (cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+    (cmakeFeature "FETCHCONTENT_TRY_FIND_PACKAGE_MODE" "ALWAYS")
+    (cmakeBool "CUDNN_FRONTEND_BUILD_SAMPLES" finalAttrs.doCheck)
+    (cmakeBool "CUDNN_FRONTEND_BUILD_TESTS" finalAttrs.doCheck)
+    (cmakeBool "CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS" false)
+  ];
+
+  checkInputs = [
+    cudnn
+    cuda_nvrtc
+    catch2_3
+    libcublas
+  ];
+
+  enableParallelBuilding = true;
+
+  propagatedBuildInputs = [
+    nlohmann_json
+  ];
+
+  doCheck = true;
+
+  postInstall = optionalString finalAttrs.doCheck ''
+    moveToOutput "bin/legacy_samples" "$legacy_samples"
+    moveToOutput "bin/samples" "$samples"
+    moveToOutput "bin/tests" "$tests"
+    if [[ -e "$out/bin" ]]
+    then
+      nixErrorLog "The bin directory in \$out should no longer exist."
+      exit 1
+    fi
+  '';
+
+  passthru.updateScript = gitUpdater {
+    inherit (finalAttrs) pname version;
+    rev-prefix = "v";
+  };
+
+  meta = with lib; {
+    description = "A c++ wrapper for the cudnn backend API";
+    homepage = "https://github.com/NVIDIA/cudnn-frontend";
+    license = licenses.mit;
+    platforms = [
+      "aarch64-linux"
+      "x86_64-linux"
+    ];
+    maintainers = (with maintainers; [ connorbaker ]) ++ teams.cuda.members;
+  };
+})

pkgs/top-level/cuda-packages.nix

@@ -79,6 +79,7 @@ let
     cudatoolkit = final.callPackage ../development/cuda-modules/cudatoolkit/redist-wrapper.nix { };
     cudatoolkit-legacy-runfile = final.callPackage ../development/cuda-modules/cudatoolkit { };
 
+    cudnn-frontend = final.callPackage ../development/cuda-modules/cudnn-frontend/default.nix { };
     saxpy = final.callPackage ../development/cuda-modules/saxpy { };
     nccl = final.callPackage ../development/cuda-modules/nccl { };
     nccl-tests = final.callPackage ../development/cuda-modules/nccl-tests { };