rarias/jungle
rarias/jungle
2
0
Fork 2
Code Issues 54 Pull Requests 7 Actions Packages Projects Releases Wiki Activity

Add oneMath

Browse Source
This commit is contained in:
Aleix Boné 2025-05-14 17:57:39 +02:00
parent 344579386d
commit 0e667db329
No known key found for this signature in database
3 changed files with 86 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
overlay.nix
View File

@ -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

1
pkgs/intel-oneapi/patch_intel.nix
View File

@ -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
'';

84
pkgs/onemath/default.nix Normal file
View File

@ -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
];
}