pyrox.dev
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 rocmUpdateScript,
6 cmake,
7 clr,
8 python3,
9 rocm-cmake,
10 sqlite,
11 boost,
12 fftw,
13 fftwFloat,
14 gtest,
15 openmp,
16 rocrand,
17 gpuTargets ? clr.localGpuTargets or clr.gpuTargets,
18}:
19
20stdenv.mkDerivation (finalAttrs: {
21 pname = "rocfft${clr.gpuArchSuffix}";
22 version = "6.3.3";
23
24 src = fetchFromGitHub {
25 owner = "ROCm";
26 repo = "rocFFT";
27 rev = "rocm-${finalAttrs.version}";
28 hash = "sha256-RrxdwZ64uC7lQzyJI1eGHX2dmRnW8TfNThnuvuz5XWo=";
29 };
30
31 nativeBuildInputs = [
32 cmake
33 clr
34 python3
35 rocm-cmake
36 ];
37
38 # FIXME: rocfft_aot_helper runs at the end of the build and has a risk of timing it out
39 # due to a long period with no terminal output
40 buildInputs = [ sqlite ];
41
42 cmakeFlags = [
43 "-DCMAKE_C_COMPILER=hipcc"
44 "-DCMAKE_CXX_COMPILER=hipcc"
45 "-DSQLITE_USE_SYSTEM_PACKAGE=ON"
46 # Manually define CMAKE_INSTALL_<DIR>
47 # See: https://github.com/NixOS/nixpkgs/pull/197838
48 "-DCMAKE_INSTALL_BINDIR=bin"
49 "-DCMAKE_INSTALL_LIBDIR=lib"
50 "-DCMAKE_INSTALL_INCLUDEDIR=include"
51 ]
52 ++ lib.optionals (gpuTargets != [ ]) [
53 "-DAMDGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
54 ];
55
56 passthru = {
57 test = stdenv.mkDerivation {
58 pname = "${finalAttrs.pname}-test";
59 inherit (finalAttrs) version src;
60
61 sourceRoot = "${finalAttrs.src.name}/clients/tests";
62
63 nativeBuildInputs = [
64 cmake
65 clr
66 rocm-cmake
67 ];
68
69 buildInputs = [
70 boost
71 fftw
72 fftwFloat
73 finalAttrs.finalPackage
74 gtest
75 openmp
76 rocrand
77 ];
78
79 cmakeFlags = [
80 "-DCMAKE_C_COMPILER=hipcc"
81 "-DCMAKE_CXX_COMPILER=hipcc"
82 ];
83
84 postInstall = ''
85 rm -r "$out/lib/fftw"
86 rmdir "$out/lib"
87 '';
88 };
89
90 benchmark = stdenv.mkDerivation {
91 pname = "${finalAttrs.pname}-benchmark";
92 inherit (finalAttrs) version src;
93
94 sourceRoot = "${finalAttrs.src.name}/clients/rider";
95
96 nativeBuildInputs = [
97 cmake
98 clr
99 rocm-cmake
100 ];
101
102 buildInputs = [
103 boost
104 finalAttrs.finalPackage
105 openmp
106 (python3.withPackages (
107 ps: with ps; [
108 pandas
109 scipy
110 ]
111 ))
112 rocrand
113 ];
114
115 cmakeFlags = [
116 "-DCMAKE_C_COMPILER=hipcc"
117 "-DCMAKE_CXX_COMPILER=hipcc"
118 ];
119
120 postInstall = ''
121 cp -a ../../../scripts/perf "$out/bin"
122 '';
123 };
124
125 samples = stdenv.mkDerivation {
126 pname = "${finalAttrs.pname}-samples";
127 inherit (finalAttrs) version src;
128
129 sourceRoot = "${finalAttrs.src.name}/clients/samples";
130
131 nativeBuildInputs = [
132 cmake
133 clr
134 rocm-cmake
135 ];
136
137 buildInputs = [
138 boost
139 finalAttrs.finalPackage
140 openmp
141 rocrand
142 ];
143
144 cmakeFlags = [
145 "-DCMAKE_C_COMPILER=hipcc"
146 "-DCMAKE_CXX_COMPILER=hipcc"
147 ];
148
149 installPhase = ''
150 runHook preInstall
151 mkdir "$out"
152 cp -a bin "$out"
153 runHook postInstall
154 '';
155 };
156
157 updateScript = rocmUpdateScript {
158 name = finalAttrs.pname;
159 inherit (finalAttrs.src) owner;
160 inherit (finalAttrs.src) repo;
161 };
162 };
163
164 requiredSystemFeatures = [ "big-parallel" ];
165
166 meta = with lib; {
167 description = "FFT implementation for ROCm";
168 homepage = "https://github.com/ROCm/rocFFT";
169 license = with licenses; [ mit ];
170 teams = [ teams.rocm ];
171 platforms = platforms.linux;
172 };
173})