From 0e667db3293fa831e92940ac65bd1f7f3d577a45 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

---
 overlay.nix                       |  1 +
 pkgs/intel-oneapi/patch_intel.nix |  1 +
 pkgs/onemath/default.nix          | 84 +++++++++++++++++++++++++++++++
 3 files changed, 86 insertions(+)
 create mode 100644 pkgs/onemath/default.nix

diff --git a/overlay.nix b/overlay.nix
index 4e8d61a7..71ee69ca 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -23,6 +23,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 { };
     # FIXME: Extend this to all linuxPackages variants. Open problem, see:
     # https://discourse.nixos.org/t/whats-the-right-way-to-make-a-custom-kernel-module-available/4636
diff --git a/pkgs/intel-oneapi/patch_intel.nix b/pkgs/intel-oneapi/patch_intel.nix
index ae43fd0a..ef95d550 100644
--- a/pkgs/intel-oneapi/patch_intel.nix
+++ b/pkgs/intel-oneapi/patch_intel.nix
@@ -149,6 +149,7 @@ stdenvNoCC.mkDerivation (finalAttrs: {
 
         sed -i 's/.*isCxx=0/isCxx=1/' $out/bin/icpx
 
+        # oneMath looks for sycl libraries in bin/../lib
         ln -s ${finalAttrs.finalPackage.pkgs.compiler}/lib $out/lib
         ln -s ${finalAttrs.finalPackage.pkgs.compiler}/include $out/include
       '';
diff --git a/pkgs/onemath/default.nix b/pkgs/onemath/default.nix
new file mode 100644
index 00000000..cbfded99
--- /dev/null
+++ b/pkgs/onemath/default.nix
@@ -0,0 +1,84 @@
+{ lib
+, fetchFromGitHub
+, cmake
+, withCFlags
+, intelPackages
+
+, mklSupport ? true
+
+, config
+
+, cudaSupport ? config.cudaSupport
+, cudaPackages ? { }
+
+, rocmSupport ? config.rocmSupport
+, hipTargets ? null # only one target at a time supported
+, rocmPackages ? { }
+}:
+
+let
+  # rocmSupport is not enough, we need a specific target
+  enableHip = rocmSupport && hipTargets != null;
+
+  stdenv = withCFlags (lib.optionals cudaSupport [ "--cuda-path=${cudaPackages.cudatoolkit}" ]) intelPackages.stdenv;
+in
+
+  # at least one backend has to be enabled
+  assert mklSupport || cudaSupport || enableHip;
+
+stdenv.mkDerivation rec {
+  pname = "oneMath";
+  version = "0.8";
+
+  src = fetchFromGitHub {
+    owner = "uxlfoundation";
+    repo = "oneMath";
+    rev = "v${version}";
+    sha256 = "sha256-xK8lKI3oqKlx3xtvdScpMq+HXAuoYCP0BZdkEqnJP5o=";
+  };
+
+  cmakeFlags = [
+    (lib.cmakeBool "ENABLE_MKLCPU_BACKEND"    mklSupport)
+    (lib.cmakeBool "ENABLE_MKLGPU_BACKEND"    mklSupport)
+
+    (lib.cmakeBool "ENABLE_CUBLAS_BACKEND"    cudaSupport)
+    (lib.cmakeBool "ENABLE_CUFFT_BACKEND"     cudaSupport)
+    (lib.cmakeBool "ENABLE_CURAND_BACKEND"    cudaSupport)
+    (lib.cmakeBool "ENABLE_CUSOLVER_BACKEND"  cudaSupport)
+    (lib.cmakeBool "ENABLE_CUSPARSE_BACKEND"  cudaSupport)
+
+    (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" hipTargets)
+  ];
+
+  nativeBuildInputs = [ cmake ];
+  buildInputs = lib.optionals (mklSupport) [
+    intelPackages.mkl
+    intelPackages.tbb
+  ] ++ lib.optionals (enableHip) [
+    rocmPackages.rocmPath
+    rocmPackages.rocblas
+    rocmPackages.rocfft
+    rocmPackages.rocsolver
+    rocmPackages.rocrand
+    rocmPackages.rocsparse
+  ] ++ lib.optionals (cudaSupport) [
+    (lib.getDev cudaPackages.cuda_cudart)
+    cudaPackages.cudatoolkit
+
+    cudaPackages.libcublas
+    cudaPackages.libcurand
+    cudaPackages.libcufft
+    cudaPackages.libcusparse
+    cudaPackages.libcusolver
+  ];
+
+}