From ad0b03172645dc68f0d2ad118dc9b4c70ff82a14 Mon Sep 17 00:00:00 2001
From: Conroy Cheers <conroy@corncheese.org>
Date: Tue, 22 Oct 2024 14:18:02 +1100
Subject: [PATCH] onnxruntime: add ncclSupport flag

---
 pkgs/development/libraries/onnxruntime/default.nix | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/libraries/onnxruntime/default.nix b/pkgs/development/libraries/onnxruntime/default.nix
index 9ac43cc4b1d3..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:
 
@@ -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 [