forked from rarias/jungle
85 lines
2.2 KiB
Nix
85 lines
2.2 KiB
Nix
{ 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
|
|
];
|
|
|
|
}
|