tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
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 hiprand,
18 gpuTargets ? clr.localGpuTargets or clr.gpuTargets,
19}:
20
21stdenv.mkDerivation (finalAttrs: {
22 pname = "rocfft${clr.gpuArchSuffix}";
23 version = "7.1.1";
24
25 src = fetchFromGitHub {
26 owner = "ROCm";
27 repo = "rocFFT";
28 rev = "rocm-${finalAttrs.version}";
29 hash = "sha256-1Ho3b5NmnzfLhDnvH6FECigs6OgpbJrxw4EnqKyaHA0=";
30 };
31
32 nativeBuildInputs = [
33 cmake
34 clr
35 python3
36 rocm-cmake
37 ];
38
39 buildInputs = [
40 sqlite
41 hiprand
42 ];
43
44 patches = [
45 # Fixes build timeout due to no log output during rocfft_aot step
46 ./log-every-n-aot-jobs.patch
47 ];
48
49 cmakeFlags = [
50 "-DSQLITE_USE_SYSTEM_PACKAGE=ON"
51 "-DHIP_PLATFORM=amd"
52 "-DBUILD_CLIENTS=OFF"
53 "-DBUILD_SHARED_LIBS=ON"
54 "-DUSE_HIPRAND=ON"
55 "-DROCFFT_KERNEL_CACHE_ENABLE=ON"
56 # Manually define CMAKE_INSTALL_<DIR>
57 # See: https://github.com/NixOS/nixpkgs/pull/197838
58 "-DCMAKE_INSTALL_BINDIR=bin"
59 "-DCMAKE_INSTALL_LIBDIR=lib"
60 "-DCMAKE_INSTALL_INCLUDEDIR=include"
61 ]
62 ++ lib.optionals (gpuTargets != [ ]) [
63 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
64 ];
65
66 passthru = {
67 test = stdenv.mkDerivation {
68 pname = "${finalAttrs.pname}-test";
69 inherit (finalAttrs) version src;
70
71 sourceRoot = "${finalAttrs.src.name}/clients/tests";
72
73 nativeBuildInputs = [
74 cmake
75 clr
76 rocm-cmake
77 ];
78
79 buildInputs = [
80 boost
81 fftw
82 fftwFloat
83 finalAttrs.finalPackage
84 gtest
85 openmp
86 rocrand
87 hiprand
88 ];
89
90 postInstall = ''
91 rm -r "$out/lib/fftw"
92 rmdir "$out/lib"
93 '';
94 };
95
96 benchmark = stdenv.mkDerivation {
97 pname = "${finalAttrs.pname}-benchmark";
98 inherit (finalAttrs) version src;
99
100 sourceRoot = "${finalAttrs.src.name}/clients/rider";
101
102 nativeBuildInputs = [
103 cmake
104 clr
105 rocm-cmake
106 ];
107
108 buildInputs = [
109 boost
110 finalAttrs.finalPackage
111 openmp
112 (python3.withPackages (
113 ps: with ps; [
114 pandas
115 scipy
116 ]
117 ))
118 rocrand
119 ];
120
121 postInstall = ''
122 cp -a ../../../scripts/perf "$out/bin"
123 '';
124 };
125
126 samples = stdenv.mkDerivation {
127 pname = "${finalAttrs.pname}-samples";
128 inherit (finalAttrs) version src;
129
130 sourceRoot = "${finalAttrs.src.name}/clients/samples";
131
132 nativeBuildInputs = [
133 cmake
134 clr
135 rocm-cmake
136 ];
137
138 buildInputs = [
139 boost
140 finalAttrs.finalPackage
141 openmp
142 rocrand
143 ];
144
145 installPhase = ''
146 runHook preInstall
147 mkdir "$out"
148 cp -a bin "$out"
149 runHook postInstall
150 '';
151 };
152
153 updateScript = rocmUpdateScript {
154 name = finalAttrs.pname;
155 inherit (finalAttrs.src) owner;
156 inherit (finalAttrs.src) repo;
157 };
158 };
159
160 requiredSystemFeatures = [ "big-parallel" ];
161
162 meta = {
163 description = "FFT implementation for ROCm";
164 homepage = "https://github.com/ROCm/rocFFT";
165 license = with lib.licenses; [ mit ];
166 teams = [ lib.teams.rocm ];
167 platforms = lib.platforms.linux;
168 };
169})