tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 stdenv,
3 lib,
4 fetchFromGitHub,
5 cmake,
6 mpi,
7 blas,
8 gfortran,
9 llvmPackages,
10 cudaPackages,
11 rocmPackages,
12 config,
13 gpuBackend ? (
14 if config.cudaSupport then
15 "cuda"
16 else if config.rocmSupport then
17 "rocm"
18 else
19 "none"
20 ),
21}:
22
23assert builtins.elem gpuBackend [
24 "none"
25 "cuda"
26 "rocm"
27];
28
29stdenv.mkDerivation (finalAttrs: {
30 pname = "spla";
31 version = "1.6.1";
32
33 src = fetchFromGitHub {
34 owner = "eth-cscs";
35 repo = "spla";
36 rev = "v${finalAttrs.version}";
37 hash = "sha256-fNH1IOKV1Re8G7GH9Xfn3itR80eonTbEGKQRRD16/2k=";
38 };
39
40 outputs = [
41 "out"
42 "dev"
43 ];
44
45 postPatch = ''
46 substituteInPlace src/gpu_util/gpu_blas_api.hpp \
47 --replace '#include <rocblas.h>' '#include <rocblas/rocblas.h>'
48 '';
49
50 nativeBuildInputs = [
51 cmake
52 gfortran
53 ];
54
55 buildInputs = [
56 blas
57 mpi
58 ]
59 ++ lib.optional (gpuBackend == "cuda") cudaPackages.cudatoolkit
60 ++ lib.optionals (gpuBackend == "rocm") [
61 rocmPackages.clr
62 rocmPackages.rocblas
63 ]
64 ++ lib.optional stdenv.hostPlatform.isDarwin llvmPackages.openmp;
65
66 cmakeFlags = [
67 "-DSPLA_OMP=ON"
68 "-DSPLA_FORTRAN=ON"
69 "-DSPLA_INSTALL=ON"
70 # Required due to broken CMake files
71 "-DCMAKE_INSTALL_LIBDIR=lib"
72 "-DCMAKE_INSTALL_INCLUDEDIR=include"
73 ]
74 ++ lib.optional (gpuBackend == "cuda") "-DSPLA_GPU_BACKEND=CUDA"
75 ++ lib.optional (gpuBackend == "rocm") [ "-DSPLA_GPU_BACKEND=ROCM" ];
76
77 preFixup = ''
78 substituteInPlace $out/lib/cmake/SPLA/SPLASharedTargets-release.cmake \
79 --replace-fail "\''${_IMPORT_PREFIX}" "$out"
80 '';
81
82 meta = {
83 description = "Specialized Parallel Linear Algebra, providing distributed GEMM functionality for specific matrix distributions with optional GPU acceleration";
84 homepage = "https://github.com/eth-cscs/spla";
85 license = lib.licenses.bsd3;
86 maintainers = [ lib.maintainers.sheepforce ];
87 };
88})