From 4c1791c95d64062d06bc032275f75d7e48725ff2 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 22cc2a7..fed652b 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -45,6 +45,7 @@ let
     #nix-wrap = callPackage ./pkgs/nix-wrap/default.nix { };
     nodes = callPackage ./pkgs/nodes/default.nix { };
     nosv = callPackage ./pkgs/nosv/default.nix { };
+    oneMath = callPackage ./pkgs/onemath/default.nix { };
     openmp = callPackage ./pkgs/llvm-ompss2/openmp.nix { monorepoSrc = final.clangOmpss2Unwrapped.src; version = final.clangOmpss2Unwrapped.version; };
     openmpv = final.openmp.override { enableNosv = true; enableOvni = true; };
     osumb = callPackage ./pkgs/osu/default.nix { };
diff --git a/pkgs/intel-oneapi/patch_intel.nix b/pkgs/intel-oneapi/patch_intel.nix
index 4e74326..0987d33 100644
--- a/pkgs/intel-oneapi/patch_intel.nix
+++ b/pkgs/intel-oneapi/patch_intel.nix
@@ -157,6 +157,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 0000000..cbfded9
--- /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
+  ];
+
+}