diff options
Diffstat (limited to 'pkgs/development/libraries/onnxruntime/default.nix')
-rw-r--r-- | pkgs/development/libraries/onnxruntime/default.nix | 8 |
1 files changed, 5 insertions, 3 deletions
diff --git a/pkgs/development/libraries/onnxruntime/default.nix b/pkgs/development/libraries/onnxruntime/default.nix index bb344664e659..a1bfe9dced1d 100644 --- 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: @@ -149,7 +150,7 @@ effectiveStdenv.mkDerivation rec { numpy pybind11 packaging - ]) ++ lib.optionals effectiveStdenv.isDarwin [ + ]) ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [ Foundation libiconv ] ++ lib.optionals cudaSupport (with cudaPackages; [ @@ -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 [ |