From 39157ea318088ee59e99a1258c2474e06009c8d4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Bon=C3=A9?= <aleix.boneribo@bsc.es>
Date: Wed, 14 May 2025 17:57:39 +0200
Subject: [PATCH] Add oneMath

With support for MKL and CUDA enabled by default
---
 overlay.nix              |  1 +
 pkgs/onemath/default.nix | 73 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 pkgs/onemath/default.nix

diff --git a/overlay.nix b/overlay.nix
index 5017672..6991741 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -24,6 +24,7 @@ let
     intelPackages_2024 = final.intel-apt.hpckit_2024;
     intelPackages_2025 = final.intel-apt.hpckit_2025;
     intelPackages = final.intelPackages_2025;
+    oneMath = callPackage ./pkgs/onemath/default.nix { };
     jemallocNanos6 = callPackage ./pkgs/nanos6/jemalloc.nix { };
     lmbench = callPackage ./pkgs/lmbench/default.nix { };
     mcxx = callPackage ./pkgs/mcxx/default.nix { };
diff --git a/pkgs/onemath/default.nix b/pkgs/onemath/default.nix
new file mode 100644
index 0000000..f320f8f
--- /dev/null
+++ b/pkgs/onemath/default.nix
@@ -0,0 +1,73 @@
+{ lib
+, intelPackages
+, withCFlags
+, fetchFromGitHub
+, cmake
+, cudaPackages
+, enableNvidia ? true
+, withHipTargets ? null # only one target at a time supported
+, rocmPackages
+, enableMkl ? true
+}:
+
+let
+  enableHip = withHipTargets != null;
+
+  stdenv = withCFlags (lib.optionals enableNvidia [ "--cuda-path=${cudaPackages.cudatoolkit}" ]) intelPackages.stdenv;
+in
+stdenv.mkDerivation rec {
+  pname = "oneMath";
+  version = "0.7";
+
+  src = fetchFromGitHub {
+    owner = "uxlfoundation";
+    repo = "oneMath";
+    rev = "v${version}";
+    sha256 = "sha256-De04PUmI68Jx4rJ+MRb+RepayZCapgnouCUCrAu6G38=";
+  };
+
+  cmakeFlags = [
+    (lib.cmakeBool "ENABLE_MKLCPU_BACKEND"    enableMkl)
+    (lib.cmakeBool "ENABLE_MKLGPU_BACKEND"    enableMkl)
+
+    (lib.cmakeBool "ENABLE_CUBLAS_BACKEND"    enableNvidia)
+    (lib.cmakeBool "ENABLE_CUFFT_BACKEND"     enableNvidia)
+    (lib.cmakeBool "ENABLE_CURAND_BACKEND"    enableNvidia)
+    (lib.cmakeBool "ENABLE_CUSOLVER_BACKEND"  enableNvidia)
+    (lib.cmakeBool "ENABLE_CUSPARSE_BACKEND"  enableNvidia)
+
+    (lib.cmakeBool "ENABLE_ROCBLAS_BACKEND"   enableHip)
+    (lib.cmakeBool "ENABLE_ROCFFT_BACKEND"    enableHip)
+    (lib.cmakeBool "ENABLE_ROCSOLVER_BACKEND" enableHip)
+    (lib.cmakeBool "ENABLE_ROCRAND_BACKEND"   enableHip)
+    (lib.cmakeBool "ENABLE_ROCSPARSE_BACKEND" enableHip)
+
+    (lib.cmakeBool "BUILD_FUNCTIONAL_TESTS" false)
+    (lib.cmakeBool "BUILD_EXAMPLES" false)
+  ] ++ lib.optionals enableHip [
+    (lib.cmakeFeature "HIP_TARGETS" withHipTargets)
+  ];
+
+  nativeBuildInputs = [ cmake ];
+  buildInputs = lib.optionals (enableMkl) [
+    intelPackages.mkl
+    intelPackages.tbb
+  ] ++ lib.optionals (enableHip) [
+    rocmPackages.rocmPath
+    rocmPackages.rocblas
+    rocmPackages.rocfft
+    rocmPackages.rocsolver
+    rocmPackages.rocrand
+    rocmPackages.rocsparse
+  ] ++ lib.optionals (enableNvidia) [
+    (lib.getDev cudaPackages.cuda_cudart)
+    cudaPackages.cudatoolkit
+
+    cudaPackages.libcublas
+    cudaPackages.libcurand
+    cudaPackages.libcufft
+    cudaPackages.libcusparse
+    cudaPackages.libcusolver
+  ];
+
+}