74 lines
2.0 KiB
Nix
74 lines
2.0 KiB
Nix
{ 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
|
|
];
|
|
|
|
}
|