diff --git a/pkgs/development/libraries/science/math/tensorrt/extension.nix b/pkgs/development/libraries/science/math/tensorrt/extension.nix index b4f47a8969c1..57e71c4928a3 100644 --- a/pkgs/development/libraries/science/math/tensorrt/extension.nix +++ b/pkgs/development/libraries/science/math/tensorrt/extension.nix @@ -29,22 +29,91 @@ final: prev: let else throw "tensorrt-${tensorRTDefaultVersion} does not support your cuda version ${cudaVersion}"; }; in allBuilds // defaultBuild; + tarballURL = + {fullVersion, fileVersionCuda, fileVersionCudnn ? null} : + "TensorRT-${fullVersion}.Linux.x86_64-gnu.cuda-${fileVersionCuda}" + + lib.optionalString (fileVersionCudnn != null) ".cudnn${fileVersionCudnn}" + + ".tar.gz"; + tensorRTVersions = { + "8.6.0" = [ + rec { + fileVersionCuda = "11.8"; + fullVersion = "8.6.0.12"; + sha256 = "sha256-wXMqEJPFerefoLaH8GG+Np5EnJwXeStmDzZj7Nj6e2M="; + tarball = tarballURL { inherit fileVersionCuda fullVersion; }; + supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ]; + } + ]; + "8.5.3" = [ + rec { + fileVersionCuda = "11.8"; + fileVersionCudnn = "8.6"; + fullVersion = "8.5.3.1"; + sha256 = "sha256-BNeuOYvPTUAfGxI0DVsNrX6Z/FAB28+SE0ptuGu7YDY="; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ]; + } + rec { + fileVersionCuda = "10.2"; + fileVersionCudnn = "8.6"; + fullVersion = "8.5.3.1"; + sha256 = "sha256-WCt6yfOmFbrjqdYCj6AE2+s2uFpISwk6urP+2I0BnGQ="; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "10.2" ]; + } + ]; + "8.5.2" = [ + rec { + fileVersionCuda = "11.8"; + fileVersionCudnn = "8.6"; + fullVersion = "8.5.2.2"; + sha256 = "sha256-Ov5irNS/JETpEz01FIFNMs9YVmjGHL7lSXmDpgCdgao="; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ]; + } + rec { + fileVersionCuda = "10.2"; + fileVersionCudnn = "8.6"; + fullVersion = "8.5.2.2"; + sha256 = "sha256-UruwQShYcHLY5d81lKNG7XaoUsZr245c+PUpUN6pC5E="; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "10.2" ]; + } + ]; + "8.5.1" = [ + rec { + fileVersionCuda = "11.8"; + fileVersionCudnn = "8.6"; + fullVersion = "8.5.1.7"; + sha256 = "sha256-Ocx/B3BX0TY3lOj/UcTPIaXb7M8RFrACC6Da4PMGMHY="; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ]; + } + rec { + fileVersionCuda = "10.2"; + fileVersionCudnn = "8.6"; + fullVersion = "8.5.1.7"; + sha256 = "sha256-CcFGJhw7nFdPnSYYSxcto2MHK3F84nLQlJYjdIw8dPM="; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "10.2" ]; + } + ]; "8.4.0" = [ rec { fileVersionCuda = "11.6"; fileVersionCudnn = "8.3"; fullVersion = "8.4.0.6"; sha256 = "sha256-DNgHHXF/G4cK2nnOWImrPXAkOcNW6Wy+8j0LRpAH/LQ="; - tarball = "TensorRT-${fullVersion}.Linux.x86_64-gnu.cuda-${fileVersionCuda}.cudnn${fileVersionCudnn}.tar.gz"; - supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" ]; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; + supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" ]; } rec { fileVersionCuda = "10.2"; fileVersionCudnn = "8.3"; fullVersion = "8.4.0.6"; sha256 = "sha256-aCzH0ZI6BrJ0v+e5Bnm7b8mNltA7NNuIa8qRKzAQv+I="; - tarball = "TensorRT-${fullVersion}.Linux.x86_64-gnu.cuda-${fileVersionCuda}.cudnn${fileVersionCudnn}.tar.gz"; + tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; }; supportedCudaVersions = [ "10.2" ]; } ]; @@ -60,7 +129,8 @@ final: prev: let "11.4" = "8.4.0"; "11.5" = "8.4.0"; "11.6" = "8.4.0"; - "11.7" = "8.4.0"; + "11.7" = "8.5.3"; + "11.8" = "8.5.3"; }.${cudaVersion} or "8.4.0"; in tensorRTPackages diff --git a/pkgs/development/libraries/science/math/tensorrt/generic.nix b/pkgs/development/libraries/science/math/tensorrt/generic.nix index 3447087051f1..492fde77e516 100644 --- a/pkgs/development/libraries/science/math/tensorrt/generic.nix +++ b/pkgs/development/libraries/science/math/tensorrt/generic.nix @@ -9,13 +9,13 @@ }: { fullVersion -, fileVersionCudnn +, fileVersionCudnn ? null , tarball , sha256 , supportedCudaVersions ? [ ] }: -assert lib.assertMsg (lib.strings.versionAtLeast cudnn.version fileVersionCudnn) +assert fileVersionCudnn == null || lib.assertMsg (lib.strings.versionAtLeast cudnn.version fileVersionCudnn) "This version of TensorRT requires at least cuDNN ${fileVersionCudnn} (current version is ${cudnn.version})"; stdenv.mkDerivation rec { diff --git a/pkgs/games/katago/default.nix b/pkgs/games/katago/default.nix index a3d7ed5cba9d..bc3ac458de46 100644 --- a/pkgs/games/katago/default.nix +++ b/pkgs/games/katago/default.nix @@ -14,28 +14,26 @@ , openssl , writeShellScriptBin , enableAVX2 ? stdenv.hostPlatform.avx2Support +, backend ? "opencl" , enableBigBoards ? false -, enableCuda ? false , enableContrib ? false -, enableGPU ? true , enableTcmalloc ? true }: -assert !enableGPU -> ( - !enableCuda); +assert lib.assertOneOf "backend" backend [ "opencl" "cuda" "tensorrt" "eigen" ]; # N.b. older versions of cuda toolkit (e.g. 10) do not support newer versions # of gcc. If you need to use cuda10, please override stdenv with gcc8Stdenv stdenv.mkDerivation rec { pname = "katago"; - version = "1.11.0"; - githash = "d8d0cd76cf73df08af3d7061a639488ae9494419"; + version = "1.12.4"; + githash = "75280bf26582090dd4985dca62bc7124116c856d"; src = fetchFromGitHub { owner = "lightvector"; repo = "katago"; rev = "v${version}"; - sha256 = "sha256-TZKkkYe2PPzgPhItBZBSJDwU3anhsujuCGIYru55OtU="; + sha256 = "sha256-1rznAxEFJ/Ah5/WiSwc+rtITOUOPYOua5BLKeqHOBr0="; }; fakegit = writeShellScriptBin "git" "echo ${githash}"; @@ -48,13 +46,17 @@ stdenv.mkDerivation rec { buildInputs = [ libzip boost - ] ++ lib.optionals (!enableGPU) [ + ] ++ lib.optionals (backend == "eigen") [ eigen - ] ++ lib.optionals (enableGPU && enableCuda) [ + ] ++ lib.optionals (backend == "cuda") [ cudaPackages.cudnn cudaPackages.cudatoolkit mesa.drivers - ] ++ lib.optionals (enableGPU && !enableCuda) [ + ] ++ lib.optionals (backend == "tensorrt") [ + cudaPackages.cudatoolkit + cudaPackages.tensorrt + mesa.drivers + ] ++ lib.optionals (backend == "opencl") [ opencl-headers ocl-icd ] ++ lib.optionals enableContrib [ @@ -65,13 +67,15 @@ stdenv.mkDerivation rec { cmakeFlags = [ "-DNO_GIT_REVISION=ON" - ] ++ lib.optionals (!enableGPU) [ - "-DUSE_BACKEND=EIGEN" ] ++ lib.optionals enableAVX2 [ "-DUSE_AVX2=ON" - ] ++ lib.optionals (enableGPU && enableCuda) [ + ] ++ lib.optionals (backend == "eigen") [ + "-DUSE_BACKEND=EIGEN" + ] ++ lib.optionals (backend == "cuda") [ "-DUSE_BACKEND=CUDA" - ] ++ lib.optionals (enableGPU && !enableCuda) [ + ] ++ lib.optionals (backend == "tensorrt") [ + "-DUSE_BACKEND=TENSORRT" + ] ++ lib.optionals (backend == "opencl") [ "-DUSE_BACKEND=OPENCL" ] ++ lib.optionals enableContrib [ "-DBUILD_DISTRIBUTED=1" @@ -85,7 +89,7 @@ stdenv.mkDerivation rec { preConfigure = '' cd cpp/ - '' + lib.optionalString enableCuda '' + '' + lib.optionalString (backend == "cuda" || backend == "tensorrt") '' export CUDA_PATH="${cudaPackages.cudatoolkit}" export EXTRA_LDFLAGS="-L/run/opengl-driver/lib" ''; @@ -93,7 +97,7 @@ stdenv.mkDerivation rec { installPhase = '' runHook preInstall mkdir -p $out/bin; cp katago $out/bin; - '' + lib.optionalString enableCuda '' + '' + lib.optionalString (backend == "cuda" || backend == "tensorrt") '' wrapProgram $out/bin/katago \ --prefix LD_LIBRARY_PATH : "/run/opengl-driver/lib" '' + '' diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix index 27359d4df425..32343297b123 100644 --- a/pkgs/top-level/all-packages.nix +++ b/pkgs/top-level/all-packages.nix @@ -35399,11 +35399,15 @@ with pkgs; katago = callPackage ../games/katago { }; katagoWithCuda = katago.override { - enableCuda = true; + backend = "cuda"; }; katagoCPU = katago.override { - enableGPU = false; + backend = "eigen"; + }; + + katagoTensorRT = katago.override { + backend = "tensorrt"; }; klavaro = callPackage ../games/klavaro { };