pyrox.dev
lol
1{ lib
2, stdenv
3, fetchFromGitHub
4, cmake
5, rocm-cmake
6, rocm-runtime
7, rocm-device-libs
8, rocm-comgr
9, rocm-smi
10, hip
11, gtest
12, chrpath ? null
13, buildTests ? false
14}:
15
16assert buildTests -> chrpath != null;
17
18stdenv.mkDerivation rec {
19 pname = "rccl";
20 rocmVersion = "5.3.1";
21 version = "2.12.10-${rocmVersion}";
22
23 outputs = [
24 "out"
25 ] ++ lib.optionals buildTests [
26 "test"
27 ];
28
29 src = fetchFromGitHub {
30 owner = "ROCmSoftwarePlatform";
31 repo = "rccl";
32 rev = "rocm-${rocmVersion}";
33 hash = "sha256-whRXGD8oINDYhFs8+hEWKWoGNqacGlyy7xi8peA8Qsk=";
34 };
35
36 nativeBuildInputs = [
37 cmake
38 rocm-cmake
39 hip
40 ];
41
42 buildInputs = [
43 rocm-runtime
44 rocm-device-libs
45 rocm-comgr
46 rocm-smi
47 gtest
48 ] ++ lib.optionals buildTests [
49 chrpath
50 ];
51
52 cmakeFlags = [
53 "-DCMAKE_C_COMPILER=hipcc"
54 "-DCMAKE_CXX_COMPILER=hipcc"
55 # Manually define CMAKE_INSTALL_<DIR>
56 # See: https://github.com/NixOS/nixpkgs/pull/197838
57 "-DCMAKE_INSTALL_BINDIR=bin"
58 "-DCMAKE_INSTALL_LIBDIR=lib"
59 "-DCMAKE_INSTALL_INCLUDEDIR=include"
60 ] ++ lib.optionals buildTests [
61 "-DBUILD_TESTS=ON"
62 ];
63
64 # Replace the manually set parallel jobs to NIX_BUILD_CORES
65 postPatch = ''
66 substituteInPlace CMakeLists.txt \
67 --replace "8 P" "$NIX_BUILD_CORES P" \
68 --replace "8)" "$NIX_BUILD_CORES)"
69 '';
70
71 postInstall = lib.optionalString buildTests ''
72 mkdir -p $test/bin
73 mv $out/bin/* $test/bin
74 rmdir $out/bin
75 '';
76
77 meta = with lib; {
78 description = "ROCm communication collectives library";
79 homepage = "https://github.com/ROCmSoftwarePlatform/rccl";
80 license = with licenses; [ bsd2 bsd3 ];
81 maintainers = with maintainers; [ Madouura ];
82 broken = rocmVersion != hip.version;
83 };
84}