onnxruntime: add ncclSupport flag

NixOS · Oct 22, 2024 · 7f496b1 · 7f496b1
1 parent 367c912
commit 7f496b1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/pkgs/development/libraries/onnxruntime/default.nix b/pkgs/development/libraries/onnxruntime/default.nix
@@ -23,6 +23,7 @@
 , protobuf_21
 , pythonSupport ? true
 , cudaSupport ? config.cudaSupport
+, ncclSupport ? config.cudaSupport
 , cudaPackages ? {}
 }@inputs:
 
@@ -160,8 +161,9 @@ effectiveStdenv.mkDerivation rec {
  libcufft # cufft.h
  cudnn # cudnn.h
  cuda_cudart
+ ] ++ lib.optionals (cudaSupport && ncclSupport) (with cudaPackages; [
  nccl
- ]);
+ ]));
 
  nativeCheckInputs = [
  gtest
@@ -198,7 +200,7 @@ effectiveStdenv.mkDerivation rec {
  "-Donnxruntime_ENABLE_LTO=ON"
  "-Donnxruntime_USE_FULL_PROTOBUF=OFF"
  (lib.cmakeBool "onnxruntime_USE_CUDA" cudaSupport)
- (lib.cmakeBool "onnxruntime_USE_NCCL" cudaSupport)
+ (lib.cmakeBool "onnxruntime_USE_NCCL" (cudaSupport && ncclSupport))
  ] ++ lib.optionals pythonSupport [
  "-Donnxruntime_ENABLE_PYTHON=ON"
  ] ++ lib.optionals cudaSupport [