From 585af9ec70b6048639bfc492163875aed3c11e55 Mon Sep 17 00:00:00 2001
From: Jonas <git@finkler.me>
Date: Fri, 30 Aug 2024 13:37:35 +0200
Subject: [PATCH] umpire: Add cuda support

---
 pkgs/by-name/um/umpire/package.nix | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/pkgs/by-name/um/umpire/package.nix b/pkgs/by-name/um/umpire/package.nix
index 17fae99e313f..54340d3fdeca 100644
--- a/pkgs/by-name/um/umpire/package.nix
+++ b/pkgs/by-name/um/umpire/package.nix
@@ -2,8 +2,13 @@
 , lib
 , fetchFromGitHub
 , cmake
+, config
+, cudaSupport ? config.cudaSupport
+, cudaPackages ? null
 }:
 
+assert cudaSupport -> cudaPackages != null;
+
 stdenv.mkDerivation rec {
   pname = "umpire";
   version = "2024.07.0";
@@ -16,7 +21,22 @@ stdenv.mkDerivation rec {
     fetchSubmodules = true;
   };
 
-  nativeBuildInputs = [ cmake ];
+  nativeBuildInputs = [
+    cmake
+  ] ++ lib.optionals cudaSupport [
+    cudaPackages.cuda_nvcc
+  ];
+
+  buildInputs = lib.optionals cudaSupport (with cudaPackages; [
+    cudatoolkit
+    cuda_cudart
+  ]);
+
+  cmakeFlags = lib.optionals cudaSupport [
+    "-DCUDA_TOOLKIT_ROOT_DIR=${cudaPackages.cudatoolkit}"
+    "-DENABLE_CUDA=ON"
+    (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" cudaPackages.flags.cmakeCudaArchitecturesString)
+  ];
 
   meta = with lib; {
     description = "Application-focused API for memory management on NUMA & GPU architectures";